diff --git a/groupselect/algorithms/algorithm_dream.py b/groupselect/algorithms/algorithm_dream.py index 3f15588..523d18d 100644 --- a/groupselect/algorithms/algorithm_dream.py +++ b/groupselect/algorithms/algorithm_dream.py @@ -22,7 +22,7 @@ def algorithm_dream( fields: dict[int, FieldMode], groups: list[(int, int)], manuals: dict[int, int], - progress_func: None | Callable = None, + progress_func: None | Callable[[float], None] = None, seed: None | int = None, pareto_prob: float = 0.5, swap_rounds: int = 1, @@ -736,7 +736,7 @@ def allocate( cluster_tables, pareto_prob, n_swap_loops, - progress_bar, + progress_bar: None | Callable[[float], None], previous_meetings, no_cluster_agents, val_cluster, @@ -759,7 +759,7 @@ def allocate( for round_no in range(n_rounds): if progress_bar: - progress_bar(round_no + 1) + progress_bar((round_no + 1) / n_rounds) if not (isinstance(tables, int)): no_tables = tables[round_no] diff --git a/groupselect/algorithms/algorithm_hermes.py b/groupselect/algorithms/algorithm_hermes.py index d37d612..05e089f 100644 --- a/groupselect/algorithms/algorithm_hermes.py +++ b/groupselect/algorithms/algorithm_hermes.py @@ -25,7 +25,7 @@ def algorithm_hermes( fields: dict[int, FieldMode], groups: list[(int, int)], manuals: dict[int, int], - progress_func: None | Callable = None, + progress_func: None | Callable[[float], None] = None, seed: None | int = None, pareto_probs: dict[int, float] = None, swap_rounds: int = 1, diff --git a/groupselect/algorithms/algorithm_legacy.py b/groupselect/algorithms/algorithm_legacy.py index 05547fc..4f6993a 100644 --- a/groupselect/algorithms/algorithm_legacy.py +++ b/groupselect/algorithms/algorithm_legacy.py @@ -19,7 +19,7 @@ def algorithm_legacy( fields: dict[int, FieldMode], groups: list[(int, int)], manuals: dict[int, int], - progress_func: None | Callable = None, + progress_func: None | Callable[[float], None] = None, n_attempts: int = 100, seed: None | int = None, ) -> AllocatorResult: @@ -67,7 +67,7 @@ def algorithm_legacy( ] for n, ensemble in enumerate(allocation_attempts): if progress_func is not None: - progress_func(n) + progress_func(n / n_attempts) for n_gr, n_ppgr in groups: # Shuffle participant IDs. shuffle = list(range(len(participants))) diff --git a/groupselect/allocate_numpy.py b/groupselect/allocate_numpy.py index cd93612..b768760 100644 --- a/groupselect/allocate_numpy.py +++ b/groupselect/allocate_numpy.py @@ -18,10 +18,16 @@ def allocate_numpy( n_part_per_group: int | Iterable[int], manuals: None | dict[int, int] = None, algorithm: Algorithm | str = Algorithm.Legacy, - progress_func: None | Callable = None, + progress_func: None | Callable[[float], None] = None, settings: None | dict = None, ) -> AllocatorResult: - """Allocate participants into groups using the chosen algorithm.""" + """Allocate participants into groups using the chosen algorithm. + + `progress_func`, if given, is called repeatedly by the chosen algorithm + with a single float in [0.0, 1.0] indicating the fraction of work done + so far -- this is comparable across all algorithms, regardless of how + each one internally breaks its work into steps. + """ # Check arguments: participants. if not np.issubdtype(participants.dtype, np.integer): raise Exception("Argument participants must only contain integers.") diff --git a/groupselect/allocate_pandas.py b/groupselect/allocate_pandas.py index f98a4aa..fe4216e 100644 --- a/groupselect/allocate_pandas.py +++ b/groupselect/allocate_pandas.py @@ -22,11 +22,17 @@ def allocate_pandas( n_part_per_group: int | Iterable[int], manuals: None | dict[int, int] = None, algorithm: Algorithm | str = Algorithm.Legacy, - progress_func: None | Callable = None, + progress_func: None | Callable[[float], None] = None, settings: None | dict = None, return_full: bool = False, ) -> pd.DataFrame | tuple[pd.DataFrame, pd.DataFrame, AllocatorResult]: - """Allocate participants of a dataframe into groups.""" + """Allocate participants of a dataframe into groups. + + `progress_func`, if given, is called repeatedly by the chosen algorithm + with a single float in [0.0, 1.0] indicating the fraction of work done + so far -- this is comparable across all algorithms, regardless of how + each one internally breaks its work into steps. + """ # Check that the dataframe index is unique. if not participants.index.is_unique: raise Exception("Index of dataframe must be unique.") @@ -129,7 +135,7 @@ def allocate( n_part_per_group: int | Iterable[int], manuals: None | dict[int, int] = None, algorithm: Algorithm | str = Algorithm.Legacy, - progress_func: None | Callable = None, + progress_func: None | Callable[[float], None] = None, settings: None | dict = None, return_full: bool = False, ) -> pd.DataFrame: