Skip to content

Commit a835cc3

Browse files
Daniel Cohenfacebook-github-bot
Daniel Cohen
authored andcommitted
Add active arms property (facebook#3286)
Summary: I've been doing this by hand often lately Reviewed By: mgarrard Differential Revision: D68844139
1 parent d059120 commit a835cc3

File tree

2 files changed

+7
-0
lines changed

2 files changed

+7
-0
lines changed

ax/core/base_trial.py

+5
Original file line numberDiff line numberDiff line change
@@ -585,6 +585,11 @@ def abandoned_arms(self) -> list[Arm]:
585585
"""All abandoned arms, associated with this trial."""
586586
pass
587587

588+
@property
589+
def active_arms(self) -> list[Arm]:
590+
"""All non abandoned arms associated with this trial."""
591+
return [arm for arm in self.arms if arm not in self.abandoned_arms]
592+
588593
@abstractproperty
589594
def generator_runs(self) -> list[GeneratorRun]:
590595
"""All generator runs associated with this trial."""

ax/core/tests/test_batch_trial.py

+2
Original file line numberDiff line numberDiff line change
@@ -413,6 +413,8 @@ def test_AbandonArm(self) -> None:
413413
metadata = self.batch.abandoned_arms_metadata[0]
414414
self.assertEqual(metadata.reason, reason)
415415
self.assertEqual(metadata.name, arm.name)
416+
self.assertEqual(len(self.batch.active_arms), len(self.arms) - 1)
417+
self.assertNotIn(arm, self.batch.active_arms)
416418

417419
# Fail to abandon arm not in BatchTrial
418420
with self.assertRaises(ValueError):

0 commit comments

Comments
 (0)