Skip to content

Commit ba6b36f

Browse files
committed
[BUG] Fix MCTS.run() docstring and add __repr__ for inspectability
- Fix run() docstring: document all 3 return dict keys (candidate, sequence, score) instead of just 2 - Fix verbose param description: uses logger.debug, not print - Add __repr__ showing depth, n_iterations, n_states, experiment - Add 3 new tests for return dict keys and repr
1 parent dc1dd8d commit ba6b36f

2 files changed

Lines changed: 44 additions & 4 deletions

File tree

pyaptamer/mcts/_algorithm.py

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ class MCTS(BaseObject):
7070
>>> experiment = AptamerEvalAptaTrans(target, model, device, prot_words)
7171
>>> mcts = MCTS(depth=5, n_iterations=2, experiment=experiment)
7272
>>> candidate = mcts.run(verbose=False)
73+
>>> candidate["candidate"] # the reconstructed aptamer string # doctest: +SKIP
7374
"""
7475

7576
def __init__(
@@ -112,6 +113,14 @@ def __init__(
112113
self.base = ""
113114
self.candidate = ""
114115

116+
def __repr__(self) -> str:
117+
"""Return a human-readable representation of the MCTS configuration."""
118+
return (
119+
f"MCTS(depth={self.depth}, n_iterations={self.n_iterations}, "
120+
f"n_states={len(self.states)}, "
121+
f"experiment={self.experiment.__class__.__name__ if self.experiment else None})"
122+
)
123+
115124
def _reset(self) -> None:
116125
"""Reset the MCTS algorithm to its initial state."""
117126
self.root = TreeNode(
@@ -281,18 +290,24 @@ def _find_best_subsequence(self) -> str:
281290
def run(self, verbose: bool = True) -> dict:
282291
"""
283292
Perform a full recommendation run consisting of `self.n_iterations` rounds of
284-
(selection -> expansion -> simulation -> backpropagation)
293+
(selection -> expansion -> simulation -> backpropagation).
285294
286295
Parameters
287296
----------
288297
verbose : bool, optional, default=True
289-
Whether to print progress information.
298+
Whether to log progress information via ``logger.debug``.
299+
Messages are only emitted when the logger level is ``DEBUG``.
290300
291301
Returns
292302
-------
293303
dict
294-
Dictionary containing the final candidate sequence (`candidate`) and its
295-
score (`score`).
304+
A dictionary with the following keys:
305+
306+
- ``"candidate"`` (*str*) — The reconstructed aptamer sequence.
307+
- ``"sequence"`` (*str*) — The raw encoded sequence with direction
308+
markers (underscores), before reconstruction.
309+
- ``"score"`` (*float*) — The score of the candidate as evaluated
310+
by ``self.experiment``.
296311
"""
297312
self._reset()
298313

pyaptamer/mcts/tests/test_mcts.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -360,3 +360,28 @@ def test_run_verbose(self, mcts):
360360
# length of sequence should be 2 * 5 (i.e., 2 * 5) as it still contains
361361
# the underscores
362362
assert len(candidate["sequence"]) == 10
363+
364+
def test_run_return_dict_keys(self, mcts):
365+
"""Verify that run() returns exactly the three documented keys."""
366+
result = mcts.run(verbose=False)
367+
assert set(result.keys()) == {"candidate", "sequence", "score"}
368+
assert isinstance(result["candidate"], str)
369+
assert isinstance(result["sequence"], str)
370+
371+
def test_repr(self):
372+
"""Verify __repr__ includes configuration details."""
373+
mcts = MCTS(depth=10, n_iterations=500)
374+
r = repr(mcts)
375+
assert "MCTS(" in r
376+
assert "depth=10" in r
377+
assert "n_iterations=500" in r
378+
assert "n_states=8" in r
379+
assert "experiment=None" in r
380+
381+
def test_repr_with_experiment(self, mcts):
382+
"""Verify __repr__ shows experiment class name."""
383+
r = repr(mcts)
384+
assert "MCTS(" in r
385+
assert "experiment=" in r
386+
assert "None" not in r
387+

0 commit comments

Comments
 (0)