Skip to content

Commit

Permalink
Filter TaskParameter when using lookup callable
Browse files Browse the repository at this point in the history
  • Loading branch information
AVHopp committed Apr 3, 2024
1 parent f494309 commit 4226ab4
Showing 1 changed file with 9 additions and 1 deletion.
10 changes: 9 additions & 1 deletion baybe/simulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -607,9 +607,17 @@ def _look_up_target_values(
# column ordering, which is not robust. Instead, the callable should return
# a dataframe with properly labeled columns.

# For evaluating the lookup function, we need to filter the task parameter
task_param_names = [
p.name for p in campaign.parameters if isinstance(p, TaskParameter)
]
filtered_queries = queries.drop(columns=task_param_names)

# Since the return of a lookup function is a a tuple, the following code stores
# tuples of floats in a single column with label 0:
measured_targets = queries.apply(lambda x: lookup(*x.values), axis=1).to_frame()
measured_targets = filtered_queries.apply(
lambda x: lookup(*x.values), axis=1
).to_frame()
# We transform this column to a DataFrame in which there is an individual
# column for each of the targets....
split_target_columns = pd.DataFrame(
Expand Down

0 comments on commit 4226ab4

Please sign in to comment.