diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..e91ec1b --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,59 @@ +name: CI + +on: + pull_request: + push: + branches: + - main + +jobs: + # 1) Run the test suite against every supported Python version. + tests: + name: Tests + runs-on: ubuntu-latest + + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12", "3.13"] + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + uv sync --all-extras + + - name: Run tests + run: | + uv run pytest -v + + # 2) Lint and formatting checks with ruff. + lint: + name: Lint & format (ruff) + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - name: Install dependencies + run: | + uv sync --group dev + + - name: Ruff lint + run: | + uv run ruff check . + + - name: Ruff format check + run: | + uv run ruff format --check . diff --git a/groupselect/__init__.py b/groupselect/__init__.py index f76d1ab..038cd93 100644 --- a/groupselect/__init__.py +++ b/groupselect/__init__.py @@ -1,7 +1,13 @@ +"""GroupSelect: allocate participants into diverse groups.""" + from groupselect.algorithms import Algorithm from groupselect.field_mode import FieldMode from groupselect.allocate_numpy import allocate_numpy -from groupselect.allocation import Allocation, AllocationEnsemble, AllocatorResult +from groupselect.allocation import ( + Allocation, + AllocationEnsemble, + AllocatorResult, +) try: diff --git a/groupselect/algorithms/__init__.py b/groupselect/algorithms/__init__.py index 0507eec..f02e0ad 100644 --- a/groupselect/algorithms/__init__.py +++ b/groupselect/algorithms/__init__.py @@ -1,3 +1,5 @@ +"""Available allocation algorithms and lookup of their implementations.""" + from enum import Enum, unique from importlib import import_module @@ -7,14 +9,20 @@ # Define Enum data type for different algorithms @unique class Algorithm(Enum): + """Identifier for a supported allocation algorithm.""" + Legacy = 0 - # Heuristic = 1 # TODO: Add heuristic algorithm + DREAM = 1 + HERMES = 2 + # Define functions of algorithms. algorithm_funcs = { algorithm: getattr( - import_module('groupselect.algorithms.algorithm_' + algorithm.name.lower()), - 'algorithm_' + algorithm.name.lower() + import_module( + "groupselect.algorithms.algorithm_" + algorithm.name.lower() + ), + "algorithm_" + algorithm.name.lower(), ) for algorithm in Algorithm } diff --git a/groupselect/algorithms/algorithm_dream.py b/groupselect/algorithms/algorithm_dream.py new file mode 100644 index 0000000..3f15588 --- /dev/null +++ b/groupselect/algorithms/algorithm_dream.py @@ -0,0 +1,813 @@ +"""DREAM allocation algorithm.""" + +import copy +import math +from itertools import product +from typing import Callable + +import numpy as np + + +from groupselect.field_mode import FieldMode +from groupselect.allocation import ( + Allocation, + ParticipantGroup, + AllocationEnsemble, + AllocatorResult, +) + + +def algorithm_dream( + participants: np.ndarray[int], + fields: dict[int, FieldMode], + groups: list[(int, int)], + manuals: dict[int, int], + progress_func: None | Callable = None, + seed: None | int = None, + pareto_prob: float = 0.5, + swap_rounds: int = 1, + cluster_tables: int = 2, +): + """Allocate participants into groups using the DREAM algorithm.""" + if not (isinstance(pareto_prob, float) and 0.0 <= pareto_prob <= 0.5): + raise Exception( + f"Pareto probability must be float in range [0.0, 0.5] but " + f"found: {pareto_prob}" + ) + + nallocations = len(groups) + progress_bar = progress_func + tables = groups[0][0] + for value in groups: + tables = value[0] + seats = value + + order_cluster = [k for k, v in fields.items() if v == FieldMode.Cluster] + order_diverse = [k for k, v in fields.items() if v == FieldMode.Diversify] + + X = participants[:, order_cluster] + X = np.unique(X) + + if len(order_cluster) >= 1: + val_cluster = X[len(X) - 1] + else: + val_cluster = "" + + Y = participants[:, order_diverse] + + lister = [None] * len(order_diverse) + + for i in range(0, len(order_diverse)): + lister[i] = [int(item[i]) for item in Y] + + lister[i] = np.unique(lister[i]) + lister[i] = lister[i].tolist() + + order_cluster_dict = dict(zip(order_cluster, [list(X)])) + order_diverse_dict = dict(zip(order_diverse, lister)) + + m_data = participants.shape[0] + + seats = math.ceil(m_data / tables) + previous_meetings = {} + + try: + random = np.random.default_rng(seed) + except: + raise Exception( + "Error: Random seed incorrect!", + "There was a problem setting the random seed. Please check " + "your input!", + ) + + if nallocations < 1: + raise Exception( + "Error: Wrong allocation number!", + "The number of computed allocations must at least be 1!", + ) + + peopledata_vals_used = [{} for i in range(m_data)] + + for i in range(m_data): + for j in order_cluster + order_diverse: + peopledata_vals_used[i][j] = int(participants[i][j]) + + if not order_diverse_dict: + raise Exception( + "Error: One diversification field required!", + "You have to set at least one field that is used to " + "diversify people across groups.", + ) + + if len(order_cluster_dict) > 1: + raise Exception( + "Error: Only one cluster field permitted. Please reduce " + "the number of cluster fields." + ) + + no_cluster_agents = 0 + + if len(order_cluster_dict) == 1: + cluster_key = next(iter(order_cluster_dict)) + no_cluster_agents = sum( + 1 + for person in peopledata_vals_used + if person[cluster_key] == val_cluster + ) + + n_swap_loops = int(swap_rounds) + if n_swap_loops < 1: + raise Exception( + "Error: at least one round of meeting optimization must be " + "specified (in *advanced settings*)" + ) + + n_results = allocate( + tables, + peopledata_vals_used, + order_cluster_dict, + order_diverse_dict, + m_data, + nallocations, + cluster_tables, + pareto_prob, + n_swap_loops, + progress_bar, + previous_meetings, + no_cluster_agents, + val_cluster, + manuals, + random, + ) + + allocation_results = n_results + + # Select the sample with maximum number of meetings. + final_results2 = max(allocation_results) + + return AllocatorResult(final_results2) + + +# --- Shared allocation machinery --------------------------------------- +# +# The functions below implement DREAM's core loop: build an initial random +# seating, then run rounds of Pareto swaps that trade off demographic +# balance against unique-meeting count. HERMES (algorithm_hermes.py) reuses +# all of it unchanged, overriding only `pareto_swaps` (injected into +# `run_round` below) to weigh each diversity field's swap-acceptance +# probability individually instead of using one fixed probability/threshold +# for every field. `evaluate_actions`, `evaluate_demographics`, and +# `select_key` accept optional `threshold`/`pareto_prob` overrides for the +# same reason, defaulting to DREAM's original fixed values. + + +def evaluate_meetings(table, previous_meetings): + """Sum each agent's meeting count with the other agents at its table.""" + total_meetings = {} + for i in range(len(table)): + for j in range(i + 1, len(table)): + agent1, agent2 = min(table[i], table[j]), max(table[i], table[j]) + # Sum the values from pairs_dict for the pair of agents + total_meetings[agent1] = total_meetings.get( + agent1, 0 + ) + previous_meetings.get((agent1, agent2), 0) + total_meetings[agent2] = total_meetings.get( + agent2, 0 + ) + previous_meetings.get((agent1, agent2), 0) + + return total_meetings + + +def evaluate_swap( + original_id, + swap_id, + allocations, + table_meeting_evaluations, + previous_meetings, +): + """Return the meeting-count change from swapping two agents' tables.""" + for index, table in enumerate(allocations): + if swap_id in table: + swap_table = index + if original_id in table: + table_no = index + + original_meetings = sum( + x for x in table_meeting_evaluations[table_no].values() + ) + sum(x for x in table_meeting_evaluations[swap_table].values()) + original_table = allocations[table_no] + swap_table = allocations[swap_table] + original_table_2 = [ + swap_id if x == original_id else x for x in original_table + ] + swap_table_2 = [original_id if x == swap_id else x for x in swap_table] + meetings_1 = evaluate_meetings(original_table_2, previous_meetings) + meetings_2 = evaluate_meetings(swap_table_2, previous_meetings) + + new_meetings = sum(x for x in meetings_1.values()) + sum( + x for x in meetings_2.values() + ) + + return original_meetings - new_meetings + + +def generate_combinations(demogs, info): + """Count, per category combination, how many values differ from info.""" + demographics = list(demogs.keys()) + + combinations_count = {} + + for values in product( + *[ + demogs[demographic] + [info[demographic]] + for demographic in demographics + ] + ): + combination = tuple(values) + count = len(demogs) - sum(1 for v in combination if v in info.values()) + combinations_count[combination] = count + return combinations_count + + +def calculate_ideal_balance(cats_diverse, m_data, people): + """Return each category's target share of the whole participant pool.""" + ideal_balance = {} + + i = 0 + for demog in cats_diverse: + i += 1 + counts = [0] * len(cats_diverse[demog]) + for row in people: + for i, category in enumerate(cats_diverse[demog]): + if row[demog] == category: + counts[i] += 1 + ideal_balance[demog] = [count / m_data for count in counts] + + return ideal_balance + + +def evaluate_actions( + ideal_dist, table_dist, cat_labels, table_size, threshold: float = 0.0 +): + """For each over-represented label, list which labels to swap in.""" + table_discrepancies = [y - x for y, x in zip(table_dist, ideal_dist)] + actions = {} + + for index, label in enumerate(cat_labels): + actions_for_label = [] + if table_dist[index] > ideal_dist[index]: + for a, b in zip(table_discrepancies, cat_labels): + if a < threshold: + actions_for_label.append(b) + actions[label] = actions_for_label + return actions + + +def evaluate_demographics( + temp_allocations, + table_no, + people, + cats_diverse, + m_data, + threshold_func: Callable[[int], float] = lambda demog: 0.0, +): + """Return a table's demographic imbalance and possible swap actions.""" + table = temp_allocations[table_no] + + table_data = {} + for index in table: + table_data[index] = people[index] + + ideal_balance = calculate_ideal_balance(cats_diverse, m_data, people) + + table_balance = {} + table_actions = {} + table_distances = {} + table_length = len(table) + for demog in cats_diverse: + counts = [0] * len(cats_diverse[demog]) + for person in table_data.values(): + for i, category in enumerate(cats_diverse[demog]): + if person.get(demog) == category: + counts[i] += 1 + table_balance[demog] = [count / table_length for count in counts] + + table_distances[demog] = sum( + [ + abs(x - y) + for x, y in zip(ideal_balance[demog], table_balance[demog]) + ] + ) / len(ideal_balance[demog]) + + table_actions[demog] = evaluate_actions( + ideal_balance[demog], + table_balance[demog], + cats_diverse[demog], + len(table), + threshold=threshold_func(demog), + ) + + return table_distances, table_actions + + +def select_key(pareto, meet, pareto_prob, random): + """Pick a swap candidate, favouring balance or unique meetings.""" + # `pareto_prob` is normally a fixed float (DREAM). HERMES instead + # passes a callable that resolves the probability from the specific + # candidate key chosen below, so it can weigh each diversity field + # individually. + if callable(pareto_prob): + pareto_prob = pareto_prob(list(pareto.keys())[0]) + + pareto_copy = pareto.copy() + meet_copy = meet.copy() + total_a = sum(pareto_copy.values()) + + if random.random() < pareto_prob: + if len(pareto_copy) == 1: + return next(iter(pareto_copy.keys())) + + cumulative_prob_a = {} + cumulative_sum = 0 + for key, value in pareto_copy.items(): + cumulative_sum += value / total_a + cumulative_prob_a[key] = cumulative_sum + + rand_num = random.random() + for key, prob in cumulative_prob_a.items(): + if rand_num <= prob: + return key + else: + meet_copy = { + key: value + for key, value in meet_copy.items() + if meet_copy[key] >= 0 + } + if len(meet_copy) == 0: + return None + if len(meet_copy) == 1: + return next(iter(meet_copy.keys())) + + total_b = sum(meet_copy.values()) + cumulative_prob_b = {} + cumulative_sum = 0 + for key, value in meet_copy.items(): + cumulative_sum += value / total_b + cumulative_prob_b[key] = cumulative_sum + rand_num = random.random() + + for key, prob in cumulative_prob_b.items(): + if rand_num <= prob: + return key + + +def pareto_swaps( + shuffled_pids, + cluster_individuals, + cluster_table_index, + temp_allocations, + people, + cats_diverse, + manual_pids, + previous_meetings, + m_data, + pareto_prob, + random, + threshold_func: Callable[[int], float] = lambda demog: 0.0, +): + """Select the best swap for a round using a single fixed threshold.""" + temp_allocations_update = temp_allocations.copy() + + table_meeting_evaluations = {} + table_demog_evaluations = {} + for index, table in enumerate(temp_allocations_update): + table_meeting_evaluations[index] = evaluate_meetings( + table, previous_meetings + ) + table_demog_evaluations[index] = {} + table_demog_evaluations[index] = evaluate_demographics( + temp_allocations_update, + index, + people, + cats_diverse, + m_data, + threshold_func=threshold_func, + ) + + for pid in shuffled_pids: + for index, table in enumerate(temp_allocations_update): + if pid in table: + table_no = index + + pid_info = { + key: people[pid][key] for key in people[pid] if key in cats_diverse + } + + candidate_demogs = {} + + for demog in cats_diverse: + candidate_demogs[demog] = table_demog_evaluations[table_no][1][ + demog + ][pid_info[demog]] + + candidate_profiles = generate_combinations(candidate_demogs, pid_info) + candidate_swaps = {} + + for profile in candidate_profiles: + if pid in cluster_individuals: + candidate_swap_tables = [ + x + for x in table_demog_evaluations + if (x != table_no) and (x in cluster_table_index) + ] + else: + candidate_swap_tables = [ + x for x in table_demog_evaluations if x != table_no + ] + for candidate_table in candidate_swap_tables: + pareto_score = 0 + pareto_profile = table_demog_evaluations[candidate_table][1] + table_valid = True + for index, demog in enumerate(pareto_profile): + if ( + pid_info[demog] + in pareto_profile[demog][profile[index]] + ): + pareto_score += 1 + elif pid_info[demog] != profile[index]: + table_valid = False + break + if table_valid: + if pid in cluster_individuals: + for swap_pid in temp_allocations_update[ + candidate_table + ]: + if swap_pid not in manual_pids: + if ( + tuple( + people[swap_pid][key] + for key in people[swap_pid] + if key in cats_diverse + ) + == profile + ): + candidate_swaps[swap_pid] = ( + pareto_score + + candidate_profiles[profile] + ) + else: + for swap_pid in temp_allocations_update[ + candidate_table + ]: + if swap_pid not in cluster_individuals: + if swap_pid not in manual_pids: + if ( + tuple( + people[swap_pid][key] + for key in people[swap_pid] + if key in cats_diverse + ) + == profile + ): + candidate_swaps[swap_pid] = ( + pareto_score + + candidate_profiles[profile] + ) + + if len(candidate_swaps) == 0: + continue + + candidate_meetings = {} + for swap in candidate_swaps: + candidate_meetings[swap] = evaluate_swap( + pid, + swap, + temp_allocations_update, + table_meeting_evaluations, + previous_meetings, + ) + + candidate_swaps = { + key: value + for key, value in candidate_swaps.items() + if (candidate_swaps[key] > 0) + or (candidate_swaps[key] == 0 and candidate_meetings[key] > 0) + } + if len(candidate_swaps) == 0: + continue + + distinct_candidates = {} + for distinct_value in {value for value in candidate_swaps.values()}: + distinct_keys = { + key + for key, value in candidate_swaps.items() + if value == distinct_value + } + max_meetings = max( + value + for key, value in candidate_meetings.items() + if key in distinct_keys + ) + distinct_candidates.update( + { + key: value + for key, value in candidate_swaps.items() + if (value == distinct_value) + and (candidate_meetings[key] == max_meetings) + } + ) + distinct_meetings = { + key: value + for key, value in candidate_meetings.items() + if key in distinct_candidates + } + + reverse_mapping = {} + for key, value in distinct_candidates.items(): + if value not in reverse_mapping: + reverse_mapping[value] = [] + reverse_mapping[value].append(key) + + final_candidates = {} + for value, keys in reverse_mapping.items(): + final_candidates[random.choice(keys)] = value + final_meetings = { + key: value + for key, value in distinct_meetings.items() + if key in final_candidates + } + + keys_to_remove = set() + + for key in final_meetings.keys(): + if any( + final_meetings[other_key] >= final_meetings[key] + and final_candidates[other_key] > final_candidates[key] + for other_key in final_meetings.keys() + if other_key != key + ): + keys_to_remove.add(key) + for key in keys_to_remove: + del final_meetings[key] + del final_candidates[key] + + final_swap = select_key( + final_candidates, final_meetings, pareto_prob, random + ) + if final_swap == None: + continue + + for index, table in enumerate(temp_allocations_update): + if final_swap in table: + swap_table = index + + temp_allocations_update[table_no] = [ + final_swap if x == pid else x + for x in temp_allocations_update[table_no] + ] + temp_allocations_update[swap_table] = [ + pid if x == final_swap else x + for x in temp_allocations_update[swap_table] + ] + + for index in [table_no, swap_table]: + table_meeting_evaluations[index] = evaluate_meetings( + temp_allocations_update[index], previous_meetings + ) + table_demog_evaluations[index] = {} + table_demog_evaluations[index] = evaluate_demographics( + temp_allocations_update, + index, + people, + cats_diverse, + m_data, + threshold_func=threshold_func, + ) + + return temp_allocations_update + + +def run_round( + template, + n_swap_loops, + seats, + m_data, + manual_pids, + n_cluster_tables, + cats_cluster, + cats_diverse, + people, + val_cluster, + no_tables, + previous_meetings, + pareto_prob, + random, + pareto_swaps=pareto_swaps, +) -> Allocation: + """Build one allocation from a template, then run rounds of swaps.""" + allocations = copy.deepcopy(template) + + all_pids = list(range(m_data)) + + shuffled_pids = all_pids.copy() + + random.shuffle(shuffled_pids) + + shuffled_pids = [x for x in shuffled_pids if x not in manual_pids] + + cluster_table_index = list(range(n_cluster_tables)) + + if len(cats_cluster) == 1: + cluster_individuals = [] + for index, person in enumerate(people): + if person[next(iter(cats_cluster))] == val_cluster: + cluster_individuals.append(index) + cluster_individuals = [ + x for x in cluster_individuals if x not in manual_pids + ] + + chosen_chair = 0 + + total_clustering_spaces = sum( + allocations[index].count(None) for index in cluster_table_index + ) + + if len(cluster_individuals) > total_clustering_spaces: + raise ValueError( + "Too many manual allocations to clustering tables: " + "please reduce manual allocations." + ) + for agent in cluster_individuals: + agent_assigned = 0 + while agent_assigned == 0: + table_no = chosen_chair % len(cluster_table_index) + seat_no = math.floor( + chosen_chair / len(cluster_table_index) % seats + ) + if allocations[table_no][seat_no] is None: + allocations[table_no][seat_no] = agent + agent_assigned = 1 + chosen_chair += 1 + else: + cluster_individuals = [] + + non_cluster_individuals = [ + x for x in shuffled_pids if x not in cluster_individuals + ] + chosen_chair = 0 + + for agent in non_cluster_individuals: + agent_assigned = 0 + while agent_assigned == 0: + table_no = chosen_chair % no_tables + seat_no = math.floor(chosen_chair / no_tables % seats) + if allocations[table_no][seat_no] is None: + allocations[table_no][seat_no] = agent + agent_assigned = 1 + chosen_chair += 1 + + if n_swap_loops == 1: + pareto_allocations = pareto_swaps( + shuffled_pids, + cluster_individuals, + cluster_table_index, + allocations, + people, + cats_diverse, + manual_pids, + previous_meetings, + m_data, + pareto_prob, + random, + ) + else: + pareto_allocations = pareto_swaps( + shuffled_pids, + cluster_individuals, + cluster_table_index, + allocations, + people, + cats_diverse, + manual_pids, + previous_meetings, + m_data, + pareto_prob, + random, + ) + for swap_round in range(1, n_swap_loops): + pareto_allocations = pareto_swaps( + shuffled_pids, + cluster_individuals, + cluster_table_index, + pareto_allocations, + people, + cats_diverse, + manual_pids, + previous_meetings, + m_data, + pareto_prob, + random, + ) + + for sublist in pareto_allocations: + for i in range(len(sublist)): + for j in range(i + 1, len(sublist)): + pair = ( + min(sublist[i], sublist[j]), + max(sublist[i], sublist[j]), + ) + + previous_meetings[pair] += 1 + + this_alloc = Allocation( + ParticipantGroup(list) for list in pareto_allocations + ) + + return this_alloc + + +def allocate( + tables, + peopledata_vals_used, + order_cluster_dict, + order_diverse_dict, + m_data, + nallocations, + cluster_tables, + pareto_prob, + n_swap_loops, + progress_bar, + previous_meetings, + no_cluster_agents, + val_cluster, + manuals, + random, + run_round=run_round, +): + """Build the requested number of allocation rounds and return them.""" + n_rounds = nallocations + + for i in range(m_data): + for j in range(i + 1, m_data): + pair = (i, j) + if pair not in previous_meetings: + previous_meetings[pair] = 0 + + allocation_attempts: list[AllocationEnsemble] = [ + AllocationEnsemble() for _ in range(1) + ] + + for round_no in range(n_rounds): + if progress_bar: + progress_bar(round_no + 1) + + if not (isinstance(tables, int)): + no_tables = tables[round_no] + else: + no_tables = tables + + no_larger_tables = m_data % no_tables + + seats = math.ceil(m_data / no_tables) + + min_cluster_tables = math.ceil(no_cluster_agents / seats) + + n_cluster_tables = min(min_cluster_tables + cluster_tables, tables) + no_smaller_tables = no_tables - no_larger_tables + + if no_larger_tables == 0: + template = [ + [None for s in range(seats)] for r in range(no_smaller_tables) + ] + else: + template = [ + [None for s in range(seats)] for r in range(no_larger_tables) + ] + [ + [None for s in range(seats - 1)] + for r in range(no_smaller_tables) + ] + + allocation = run_round( + template, + n_swap_loops, + seats, + m_data, + manuals, + n_cluster_tables, + order_cluster_dict, + order_diverse_dict, + peopledata_vals_used, + val_cluster, + no_tables, + previous_meetings, + pareto_prob, + random, + ) + + allocation = Allocation( + ParticipantGroup(p_id for p_id in group) for group in allocation + ) + for n, ensemble in enumerate(allocation_attempts): + ensemble.append(allocation) + + return allocation_attempts diff --git a/groupselect/algorithms/algorithm_hermes.py b/groupselect/algorithms/algorithm_hermes.py new file mode 100644 index 0000000..d37d612 --- /dev/null +++ b/groupselect/algorithms/algorithm_hermes.py @@ -0,0 +1,422 @@ +"""HERMES allocation algorithm, built on top of DREAM's swap machinery.""" + +import functools +import math +from typing import Callable + +import numpy as np + + +from groupselect.field_mode import FieldMode +from groupselect.allocation import AllocatorResult +from groupselect.algorithms.algorithm_dream import ( + allocate, + run_round, + evaluate_demographics, + select_key, + evaluate_meetings, + generate_combinations, + evaluate_swap, +) + + +def algorithm_hermes( + participants: np.ndarray[int], + fields: dict[int, FieldMode], + groups: list[(int, int)], + manuals: dict[int, int], + progress_func: None | Callable = None, + seed: None | int = None, + pareto_probs: dict[int, float] = None, + swap_rounds: int = 1, + cluster_tables: int = 2, +): + """Allocate participants into groups using the HERMES algorithm.""" + pareto_probs = pareto_probs or {} + for field_id in fields: + if fields[field_id] == FieldMode.Diversify: + if field_id not in pareto_probs: + raise Exception( + f"Algorithm HERMES requires pareto probability to be " + f"passed for diversity fields, but none found for " + f"field {field_id}." + ) + if not ( + isinstance(pareto_probs[field_id], float) + and 0.0 <= pareto_probs[field_id] <= 0.5 + ): + raise Exception( + f"Pareto probability must be float in range " + f"[0.0, 0.5] but found: {pareto_probs[field_id]}" + ) + + nallocations = len(groups) + progress_bar = progress_func + tables = groups[0][0] + for value in groups: + tables = value[0] + seats = value + + order_cluster = [k for k, v in fields.items() if v == FieldMode.Cluster] + order_diverse = [k for k, v in fields.items() if v == FieldMode.Diversify] + + X = participants[:, order_cluster] + + X = np.unique(X) + + Y = participants[:, order_diverse] + lister = [None] * len(order_diverse) + + for i in range(0, len(order_diverse)): + lister[i] = [int(item[i]) for item in Y] + + lister[i] = np.unique(lister[i]) + lister[i] = lister[i].tolist() + + if len(order_cluster) >= 1: + val_cluster = X[len(X) - 1] + else: + val_cluster = "" + + order_cluster_dict = dict(zip(order_cluster, [list(X)])) + order_diverse_dict = dict(zip(order_diverse, lister)) + + m_data = participants.shape[0] + + seats = math.ceil(m_data / tables) + previous_meetings = {} + try: + random = np.random.default_rng(seed) + except: + raise Exception( + "Error: Random seed incorrect!", + "There was a problem setting the random seed. Please check " + "your input!", + ) + + if nallocations < 1: + raise Exception( + "Error: Wrong allocation number!", + "The number of computed allocations must at least be 1!", + ) + + peopledata_vals_used = [{} for i in range(m_data)] + + for i in range(m_data): + for j in order_cluster + order_diverse: + peopledata_vals_used[i][j] = int(participants[i][j]) + + if not order_diverse_dict: + raise Exception( + "Error: One diversification field required!", + "You have to set at least one field that is used to " + "diversify people across groups.", + ) + + if len(order_cluster_dict) > 1: + raise Exception( + "Error: Only one cluster field permitted. Please reduce " + "the number of cluster fields." + ) + + no_cluster_agents = 0 + + if len(order_cluster_dict) == 1: + cluster_key = next(iter(order_cluster_dict)) + no_cluster_agents = sum( + 1 + for person in peopledata_vals_used + if person[cluster_key] == val_cluster + ) + + n_swap_loops = int(swap_rounds) + if n_swap_loops < 1: + raise Exception( + "Error: at least one round of meeting optimization must be " + "specified (in *advanced settings*)" + ) + + # HERMES reuses DREAM's allocation scaffolding (allocate -> run_round) + # unchanged, overriding only the swap-selection step (pareto_swaps + # below) so each diversity field's swap-acceptance probability is + # weighed individually instead of using one fixed probability/threshold + # for every field. + allocation_results = allocate( + tables, + peopledata_vals_used, + order_cluster_dict, + order_diverse_dict, + m_data, + nallocations, + cluster_tables, + pareto_probs, + n_swap_loops, + progress_bar, + previous_meetings, + no_cluster_agents, + val_cluster, + manuals, + random, + run_round=functools.partial(run_round, pareto_swaps=pareto_swaps), + ) + + final_results2 = max(allocation_results) + + return AllocatorResult(final_results2) + + +def pareto_swaps( + shuffled_pids, + cluster_individuals, + cluster_table_index, + temp_allocations, + people, + cats_diverse, + manual_pids, + previous_meetings, + m_data, + pareto_probs, + random, +): + """Select the best swap for a round using HERMES's field thresholds.""" + # HERMES's per-field threshold: how far a table's demographic balance + # may deviate from the panel-wide ideal before a swap is considered for + # that field, controlled by the field's own pareto probability (DREAM + # uses a fixed threshold of 0 for every field instead). Keyed directly + # by the field's own ID -- earlier code instead reverse-looked-up the + # field from its category-value list, which silently resolved to the + # wrong field whenever two diversity fields shared an identical set of + # category values. + threshold_func = lambda demog: -0.5 + pareto_probs[demog] + + temp_allocations_update = temp_allocations.copy() + + table_meeting_evaluations = {} + table_demog_evaluations = {} + for index, table in enumerate(temp_allocations_update): + table_meeting_evaluations[index] = evaluate_meetings( + table, previous_meetings + ) + table_demog_evaluations[index] = evaluate_demographics( + temp_allocations_update, + index, + people, + cats_diverse, + m_data, + threshold_func=threshold_func, + ) + + for pid in shuffled_pids: + for index, table in enumerate(temp_allocations_update): + if pid in table: + table_no = index + + pid_info = { + key: people[pid][key] for key in people[pid] if key in cats_diverse + } + + candidate_demogs = {} + + for demog in cats_diverse: + candidate_demogs[demog] = table_demog_evaluations[table_no][1][ + demog + ][pid_info[demog]] + + candidate_profiles = generate_combinations(candidate_demogs, pid_info) + candidate_swaps = {} + + # Per swap-candidate record of which diversity fields (keyed by + # their field ID) that specific swap would pareto-improve. + demog_scores: dict[int, set[int]] = {} + + for profile in candidate_profiles: + if pid in cluster_individuals: + candidate_swap_tables = [ + x + for x in table_demog_evaluations + if (x != table_no) and (x in cluster_table_index) + ] + else: + candidate_swap_tables = [ + x for x in table_demog_evaluations if x != table_no + ] + for candidate_table in candidate_swap_tables: + demog_pareto_fields = set() + pareto_score = 0 + pareto_profile = table_demog_evaluations[candidate_table][1] + table_valid = True + for index, demog in enumerate(pareto_profile): + if ( + pid_info[demog] + in pareto_profile[demog][profile[index]] + ): + pareto_score += 1 + demog_pareto_fields.add(demog) + + elif pid_info[demog] != profile[index]: + table_valid = False + break + if table_valid: + if pid in cluster_individuals: + for swap_pid in temp_allocations_update[ + candidate_table + ]: + if swap_pid not in manual_pids: + if ( + tuple( + people[swap_pid][key] + for key in people[swap_pid] + if key in cats_diverse + ) + == profile + ): + candidate_swaps[swap_pid] = ( + pareto_score + + candidate_profiles[profile] + ) + + demog_scores.setdefault( + swap_pid, set() + ).update(demog_pareto_fields) + + else: + for swap_pid in temp_allocations_update[ + candidate_table + ]: + if swap_pid not in cluster_individuals: + if swap_pid not in manual_pids: + if ( + tuple( + people[swap_pid][key] + for key in people[swap_pid] + if key in cats_diverse + ) + == profile + ): + candidate_swaps[swap_pid] = ( + pareto_score + + candidate_profiles[profile] + ) + + demog_scores.setdefault( + swap_pid, set() + ).update(demog_pareto_fields) + + if len(candidate_swaps) == 0: + continue + candidate_meetings = {} + for swap in candidate_swaps: + candidate_meetings[swap] = evaluate_swap( + pid, + swap, + temp_allocations_update, + table_meeting_evaluations, + previous_meetings, + ) + + candidate_swaps = { + key: value + for key, value in candidate_swaps.items() + if (candidate_swaps[key] > 0) + or (candidate_swaps[key] == 0 and candidate_meetings[key] > 0) + } + + if len(candidate_swaps) == 0: + continue + + distinct_candidates = {} + for distinct_value in {value for value in candidate_swaps.values()}: + distinct_keys = { + key + for key, value in candidate_swaps.items() + if value == distinct_value + } + max_meetings = max( + value + for key, value in candidate_meetings.items() + if key in distinct_keys + ) + + distinct_candidates.update( + { + key: value + for key, value in candidate_swaps.items() + if (value == distinct_value) + and (candidate_meetings[key] == max_meetings) + } + ) + distinct_meetings = { + key: value + for key, value in candidate_meetings.items() + if key in distinct_candidates + } + reverse_mapping = {} + for key, value in distinct_candidates.items(): + if value not in reverse_mapping: + reverse_mapping[value] = [] + reverse_mapping[value].append(key) + final_candidates = {} + for value, keys in reverse_mapping.items(): + final_candidates[random.choice(keys)] = value + final_meetings = { + key: value + for key, value in distinct_meetings.items() + if key in final_candidates + } + + keys_to_remove = set() + for key in final_meetings.keys(): + if any( + final_meetings[other_key] >= final_meetings[key] + and final_candidates[other_key] > final_candidates[key] + for other_key in final_meetings.keys() + if other_key != key + ): + keys_to_remove.add(key) + for key in keys_to_remove: + del final_meetings[key] + del final_candidates[key] + + # Use the highest configured probability among the diversity + # fields the chosen candidate actually pareto-improves (falling + # back to the overall maximum if it improves none, e.g. it was + # accepted purely for its meeting-uniqueness score). + def pareto_prob_for(k): + relevant = [ + pareto_probs[field_id] for field_id in demog_scores.get(k, ()) + ] + return max(relevant) if relevant else max(pareto_probs.values()) + + final_swap = select_key( + final_candidates, final_meetings, pareto_prob_for, random + ) + if final_swap == None: + continue + + for index, table in enumerate(temp_allocations_update): + if final_swap in table: + swap_table = index + + temp_allocations_update[table_no] = [ + final_swap if x == pid else x + for x in temp_allocations_update[table_no] + ] + temp_allocations_update[swap_table] = [ + pid if x == final_swap else x + for x in temp_allocations_update[swap_table] + ] + + for index in [table_no, swap_table]: + table_meeting_evaluations[index] = evaluate_meetings( + temp_allocations_update[index], previous_meetings + ) + table_demog_evaluations[index] = evaluate_demographics( + temp_allocations_update, + index, + people, + cats_diverse, + m_data, + threshold_func=threshold_func, + ) + + return temp_allocations_update diff --git a/groupselect/algorithms/algorithm_legacy.py b/groupselect/algorithms/algorithm_legacy.py index 9a97ef4..05547fc 100644 --- a/groupselect/algorithms/algorithm_legacy.py +++ b/groupselect/algorithms/algorithm_legacy.py @@ -1,47 +1,61 @@ +"""Legacy allocation algorithm.""" + from typing import Callable import numpy as np from groupselect.field_mode import FieldMode -from groupselect.allocation import Allocation, ParticipantGroup, AllocationEnsemble, AllocatorResult +from groupselect.allocation import ( + Allocation, + ParticipantGroup, + AllocationEnsemble, + AllocatorResult, +) # Legacy algorithm -def algorithm_legacy(participants: np.ndarray[int], - fields: dict[int, FieldMode], - groups: list[(int, int)], - manuals: dict[int, int], - progress_func: None | Callable = None, - n_attempts: int = 100, - seed: None | int = None) -> AllocatorResult: - +def algorithm_legacy( + participants: np.ndarray[int], + fields: dict[int, FieldMode], + groups: list[(int, int)], + manuals: dict[int, int], + progress_func: None | Callable = None, + n_attempts: int = 100, + seed: None | int = None, +) -> AllocatorResult: + """Allocate participants into groups using the legacy algorithm.""" # Check argument: seed if not (seed is None or isinstance(seed, int)): - raise Exception('Random number seed must be integer.') + raise Exception("Random number seed must be integer.") # Try to set the random seed. try: random = np.random.default_rng(seed) except: - raise Exception('Error when creating random number generator.', - 'This may be due to an incorrect seed. Please ' - 'check your input.') + raise Exception( + "Error when creating random number generator.", + "This may be due to an incorrect seed. Please check your input.", + ) # Check argument: n_attempts if not (isinstance(n_attempts, int) and n_attempts > 0): - raise Exception('Argument n_attempts must be positive integer.') + raise Exception("Argument n_attempts must be positive integer.") # Check that there is at least one diversification field defined. if not FieldMode.Diversify in fields.values(): - raise Exception('Error: One diversification field required!', - 'You have to set at least one field that is ' - 'used to diversify people across groups.') + raise Exception( + "Error: One diversification field required!", + "You have to set at least one field that is " + "used to diversify people across groups.", + ) # Reindex the field values such that the index is descending w.r.t. # the number of occurrences of the specific field value. participants = participants.copy() for field_id in fields: - field_vals, field_val_counts = np.unique(participants.T[field_id], return_counts=True) + field_vals, field_val_counts = np.unique( + participants.T[field_id], return_counts=True + ) mapping_dict = dict(zip(field_vals, field_val_counts.argsort())) mapping_func = np.vectorize(mapping_dict.get) participants.T[field_id] = mapping_func(participants.T[field_id]) @@ -49,8 +63,7 @@ def algorithm_legacy(participants: np.ndarray[int], # Generate AllocationEnsemble `n_attempts` times, where `n_attempts` # is an externally given argument. allocation_attempts: list[AllocationEnsemble] = [ - AllocationEnsemble() - for _ in range(n_attempts) + AllocationEnsemble() for _ in range(n_attempts) ] for n, ensemble in enumerate(allocation_attempts): if progress_func is not None: @@ -80,11 +93,12 @@ def algorithm_legacy(participants: np.ndarray[int], # Sample `n_allocation` allocations and repeat that `n_attempts` times. allocation_samples: list[AllocationEnsemble] = [ - AllocationEnsemble() - for _ in range(n_attempts) + AllocationEnsemble() for _ in range(n_attempts) ] for a_id in range(len(groups)): - for sample, choice in zip(allocation_samples, random.choice(n_attempts, n_attempts)): + for sample, choice in zip( + allocation_samples, random.choice(n_attempts, n_attempts) + ): sample.append(allocation_attempts[choice][a_id]) # Select the sample with maximum number of meetings. @@ -96,30 +110,34 @@ def algorithm_legacy(participants: np.ndarray[int], # Create AllocatorResult from sample max and return return AllocatorResult(ensemble=allocation_sample_max) + # Actually run the calculation for given number of participants # per groups, number of groups, and participants' data. -def _allocate_legacy_once(participants: np.ndarray[int], - fields: dict[int, FieldMode], - n_gr: int, - n_ppgr: int, - manuals: dict[int, int]) -> Allocation: +def _allocate_legacy_once( + participants: np.ndarray[int], + fields: dict[int, FieldMode], + n_gr: int, + n_ppgr: int, + manuals: dict[int, int], +) -> Allocation: # Generate empty allocations. - allocation = Allocation( - ParticipantGroup() - for _ in range(n_gr) - ) + allocation = Allocation(ParticipantGroup() for _ in range(n_gr)) # Allocate manuals for m_p_id, m_g_id in manuals.items(): # Find first empty spot allocation[m_g_id].append(m_p_id) - # Split fields into clustering and diversification fields. Order participant IDs - # by clustering and diversification field values. + # Split fields into clustering and diversification fields. Order + # participant IDs by clustering and diversification field values. fields_cluster = [k for k, v in fields.items() if v == FieldMode.Cluster] - fields_diversify = [k for k, v in fields.items() if v == FieldMode.Diversify] - p_ids_ordered = np.lexsort(participants[:, (fields_cluster + fields_diversify)[::-1]].T) + fields_diversify = [ + k for k, v in fields.items() if v == FieldMode.Diversify + ] + p_ids_ordered = np.lexsort( + participants[:, (fields_cluster + fields_diversify)[::-1]].T + ) # Loop over participants and allocate to a group. for p_id, p_details in zip(p_ids_ordered, participants[p_ids_ordered]): @@ -137,21 +155,21 @@ def _allocate_legacy_once(participants: np.ndarray[int], return allocation + # Add a participant to a group -def _allocate_person(p_id: int, - p_details: np.ndarray[int], - allocation: Allocation, - participants: np.ndarray[int], - fields_cluster: list[int], - fields_diversify: list[int], - n_ppgr: int): +def _allocate_person( + p_id: int, + p_details: np.ndarray[int], + allocation: Allocation, + participants: np.ndarray[int], + fields_cluster: list[int], + fields_diversify: list[int], + n_ppgr: int, +): # List of groups to choose from is initially list of non-full groups. groups_list = [ - g_id - for g_id, group in enumerate(allocation) - if len(group) < n_ppgr + g_id for g_id, group in enumerate(allocation) if len(group) < n_ppgr ] - # Loop over clustering fields. superior_field_filters = {} for field_id in fields_cluster: @@ -166,25 +184,31 @@ def _allocate_person(p_id: int, # Determine how often this field value occurs across groups. field_value_counts = { - g_id: _count_categories(allocation[g_id], field_id, field_val, participants) + g_id: _count_categories( + allocation[g_id], field_id, field_val, participants + ) for g_id in groups_list } - # Create a temporary list of potential groups to focus on. To begin with, this list - # contains all groups that already contain at least one participant of the same - # field value. + # Create a temporary list of potential groups to focus on. To + # begin with, this list contains all groups that already contain + # at least one participant of the same field value. groups_list_tmp = [ - g_id - for g_id in groups_list - if field_value_counts[g_id] > 0 + g_id for g_id in groups_list if field_value_counts[g_id] > 0 ] - # Then check if the temporary list of groups is enough to accommodate all participants with - # this specific field value. If not, then keep adding one more group, which is determined + # Then check if the temporary list of groups is enough to + # accommodate all participants with this specific field value. + # If not, then keep adding one more group, which is determined # from all groups from groups_list - spaces_required = _number_of_people_filtered(participants, superior_field_filters | {field_id: field_val}) + spaces_required = _number_of_people_filtered( + participants, superior_field_filters | {field_id: field_val} + ) groups_options = None - while sum((n_ppgr - len(allocation[g_id])) for g_id in groups_list_tmp) < spaces_required: + while ( + sum((n_ppgr - len(allocation[g_id])) for g_id in groups_list_tmp) + < spaces_required + ): groups_options = groups_options or sorted( [g_id for g_id in groups_list if g_id not in groups_list_tmp], key=lambda g_id: len(allocation[g_id]), @@ -197,10 +221,14 @@ def _allocate_person(p_id: int, for field_id in fields_diversify: field_val = p_details[field_id] field_value_counts = { - g_id: _count_categories(allocation[g_id], field_id, field_val, participants) + g_id: _count_categories( + allocation[g_id], field_id, field_val, participants + ) for g_id in groups_list } - field_val_counts_min = min(field_value_counts.items(), key=lambda x: x[1])[1] + field_val_counts_min = min( + field_value_counts.items(), key=lambda x: x[1] + )[1] groups_list = [ g_id for g_id in groups_list @@ -210,9 +238,7 @@ def _allocate_person(p_id: int, # Select groups with least number of participants group_size_min = min(len(allocation[g_id]) for g_id in groups_list) groups_list = [ - g_id - for g_id in groups_list - if len(allocation[g_id]) == group_size_min + g_id for g_id in groups_list if len(allocation[g_id]) == group_size_min ] # Add participant to first group from the list of groups that remain. @@ -221,9 +247,24 @@ def _allocate_person(p_id: int, return + # Count number of occurrences of field value in a participant group. -def _count_categories(group: ParticipantGroup, field_id: int, field_val: int, participants: np.ndarray[int]): - return sum(1 for p_id in group if participants[p_id][field_id] == field_val) +def _count_categories( + group: ParticipantGroup, + field_id: int, + field_val: int, + participants: np.ndarray[int], +): + return sum( + 1 for p_id in group if participants[p_id][field_id] == field_val + ) -def _number_of_people_filtered(participants: np.ndarray[int], fields: dict[int, int]): - return (participants.T[list(fields.keys())] == list(fields.values())).all(axis=0).sum() + +def _number_of_people_filtered( + participants: np.ndarray[int], fields: dict[int, int] +): + return ( + (participants.T[list(fields.keys())] == list(fields.values())) + .all(axis=0) + .sum() + ) diff --git a/groupselect/allocate_numpy.py b/groupselect/allocate_numpy.py index f7ef59e..cd93612 100644 --- a/groupselect/allocate_numpy.py +++ b/groupselect/allocate_numpy.py @@ -1,3 +1,5 @@ +"""Allocate participants given as a plain numpy array.""" + from typing import Callable, Iterable from inspect import signature from math import ceil @@ -10,74 +12,99 @@ # Generic function for allocations. -def allocate_numpy(participants: np.ndarray[int], - fields: dict[int, FieldMode], - n_part_per_group: int | Iterable[int], - manuals: None | dict[int, int] = None, - algorithm: Algorithm | str = Algorithm.Legacy, - progress_func: None | Callable = None, - settings: None | dict = None) -> AllocatorResult: - +def allocate_numpy( + participants: np.ndarray[int], + fields: dict[int, FieldMode], + n_part_per_group: int | Iterable[int], + manuals: None | dict[int, int] = None, + algorithm: Algorithm | str = Algorithm.Legacy, + progress_func: None | Callable = None, + settings: None | dict = None, +) -> AllocatorResult: + """Allocate participants into groups using the chosen algorithm.""" # Check arguments: participants. if not np.issubdtype(participants.dtype, np.integer): - raise Exception('Argument participants must only contain integers.') + raise Exception("Argument participants must only contain integers.") if not participants.ndim == 2: - raise Exception('Argument participants must be 2-dimensional array.') + raise Exception("Argument participants must be 2-dimensional array.") # Define constants from participants' data. n_participants = participants.shape[0] n_fields = participants.shape[1] # Check argument n_part_per_group - if not ((isinstance(n_part_per_group, int) and 0 < n_part_per_group <= n_participants) or - (isinstance(n_part_per_group, Iterable) and - all(isinstance(n, int) and 0 < n <= n_participants for n in n_part_per_group))): - raise Exception('Number of participants per group must be a positive integer less' - 'than the total number of participants or a list of such integers.') + if not ( + ( + isinstance(n_part_per_group, int) + and 0 < n_part_per_group <= n_participants + ) + or ( + isinstance(n_part_per_group, Iterable) + and all( + isinstance(n, int) and 0 < n <= n_participants + for n in n_part_per_group + ) + ) + ): + raise Exception( + "Number of participants per group must be a positive integer " + "less than the total number of participants or a list of " + "such integers." + ) n_part_per_group = ( [n_part_per_group] - if isinstance(n_part_per_group, int) else - list(n_part_per_group) + if isinstance(n_part_per_group, int) + else list(n_part_per_group) ) - # Compute number of groups per allocation from number of participants per group. + # Compute number of groups per allocation from number of + # participants per group. groups = [ - (ceil(n_participants / n_ppgr), n_ppgr) - for n_ppgr in n_part_per_group + (ceil(n_participants / n_ppgr), n_ppgr) for n_ppgr in n_part_per_group ] # Check arguments: manuals if manuals is not None: - if not all(isinstance(m_p_id, int) and - 0 <= m_p_id < n_participants - for m_p_id in manuals): - raise Exception('IDs of manually allocated participants must ' - 'be integers within range of number of ' - 'participants.') - if not all(isinstance(m_g_id, int) and - 0 <= m_g_id < n_gr - for n_gr, n_ppgr in groups - for m_g_id in manuals.values()): - raise Exception('Group IDs of manual allocations must be ' - 'within range of requested group sizes.') + if not all( + isinstance(m_p_id, int) and 0 <= m_p_id < n_participants + for m_p_id in manuals + ): + raise Exception( + "IDs of manually allocated participants must " + "be integers within range of number of " + "participants." + ) + if not all( + isinstance(m_g_id, int) and 0 <= m_g_id < n_gr + for n_gr, n_ppgr in groups + for m_g_id in manuals.values() + ): + raise Exception( + "Group IDs of manual allocations must be " + "within range of requested group sizes." + ) else: manuals = {} # Check arguments: fields for f_id, f_usage in fields.items(): if not (isinstance(f_id, int) and 0 <= f_id < n_fields): - raise Exception('Keys of dict containing field usage modes ' - 'must be positive integers and less than the ' - 'number of fields provided for participants.') + raise Exception( + "Keys of dict containing field usage modes " + "must be positive integers and less than the " + "number of fields provided for participants." + ) if isinstance(f_usage, str): try: - fields[f_id] = next(f for f in FieldMode - if f.name.casefold() == f_usage.casefold()) + fields[f_id] = next( + f + for f in FieldMode + if f.name.casefold() == f_usage.casefold() + ) except StopIteration: raise Exception(f"Unknown field usage type: {f_usage}") - if not all(isinstance(f_usage, FieldMode) - for f_usage in fields.values()): - raise Exception('Field usage must be valid usage type.') + if not all(isinstance(f_usage, FieldMode) for f_usage in fields.values()): + raise Exception("Field usage must be valid usage type.") for f_id in range(n_fields): if f_id not in fields: fields[f_id] = FieldMode.Ignore @@ -85,15 +112,22 @@ def allocate_numpy(participants: np.ndarray[int], # Check arguments: algorithm if isinstance(algorithm, str): try: - algorithm = next(a for a in Algorithm - if a.name.casefold() == algorithm.casefold()) + algorithm = next( + a + for a in Algorithm + if a.name.casefold() == algorithm.casefold() + ) except StopIteration: - raise Exception(f"Unknown algorithm type: {algorithm}. Please " - 'choose from: ' + ', '.join(a.name for a in Algorithm)) + raise Exception( + f"Unknown algorithm type: {algorithm}. Please " + "choose from: " + ", ".join(a.name for a in Algorithm) + ) if not (isinstance(algorithm, Algorithm)): - raise Exception('Argument algorithm must be a valid ' - 'algorithm. Please choose from: ' - + ', '.join(a.name for a in Algorithm)) + raise Exception( + "Argument algorithm must be a valid " + "algorithm. Please choose from: " + + ", ".join(a.name for a in Algorithm) + ) # Initialise empty settings dict. if settings is None: @@ -111,11 +145,7 @@ def allocate_numpy(participants: np.ndarray[int], groups=groups, manuals=manuals, progress_func=progress_func, - **{ - k: v - for k, v in settings.items() - if k in algorithm_func_args - } + **{k: v for k, v in settings.items() if k in algorithm_func_args}, ) return result diff --git a/groupselect/allocate_pandas.py b/groupselect/allocate_pandas.py index 5ef3cd3..f98a4aa 100644 --- a/groupselect/allocate_pandas.py +++ b/groupselect/allocate_pandas.py @@ -1,3 +1,5 @@ +"""Allocate participants given as a pandas DataFrame.""" + import warnings from typing import Callable, Iterable, Hashable @@ -7,48 +9,59 @@ try: import pandas as pd except ImportError: - ImportError('Package `pandas` needs to be installed in order to use the ' - 'allocate_pandas module. Please first install `pandas` via ' - 'pip or similar.') + raise ImportError( + "Package `pandas` needs to be installed in order to use " + "the allocate_pandas module. Please first install " + "`pandas` via pip or similar." + ) -def allocate_pandas(participants: pd.DataFrame, - fields: dict[Hashable, FieldMode], - n_part_per_group: int | Iterable[int], - manuals: None | dict[int, int] = None, - algorithm: Algorithm | str = Algorithm.Legacy, - progress_func: None | Callable = None, - settings: None | dict = None, - return_full: bool = False) -> pd.DataFrame | tuple[pd.DataFrame, pd.DataFrame, AllocatorResult]: +def allocate_pandas( + participants: pd.DataFrame, + fields: dict[Hashable, FieldMode], + n_part_per_group: int | Iterable[int], + manuals: None | dict[int, int] = None, + algorithm: Algorithm | str = Algorithm.Legacy, + progress_func: None | Callable = None, + settings: None | dict = None, + return_full: bool = False, +) -> pd.DataFrame | tuple[pd.DataFrame, pd.DataFrame, AllocatorResult]: + """Allocate participants of a dataframe into groups.""" # Check that the dataframe index is unique. if not participants.index.is_unique: - raise Exception('Index of dataframe must be unique.') + raise Exception("Index of dataframe must be unique.") # Check that the specified field IDs are among the dataframe columns # and are unequal to predefined names 'group' and 'allocation'. for field_id in fields: if not field_id in participants: raise Exception(f"Field not found as dataframe column: {field_id}") - if field_id in ('group', 'allocation'): - warnings.warn("The column labels 'group' and 'allocation' will be " - "overwritten in the returned allocation result.") + if field_id in ("group", "allocation"): + warnings.warn( + "The column labels 'group' and 'allocation' will be " + "overwritten in the returned allocation result." + ) # Turn fields into categories. - participants = participants.astype({ - field_id: 'category' - for field_id in fields - }) + participants = participants.astype( + {field_id: "category" for field_id in fields} + ) # Obtain codes for participants' data from pandas. - participants_codes = ( - participants[list(fields.keys())] - .apply(lambda col: col.cat.codes) + participants_codes = participants[list(fields.keys())].apply( + lambda col: col.cat.codes ) participants_numpy = participants_codes.to_numpy() fields_numpy = { participants_codes.columns.tolist().index(k): v for k, v in fields.items() } + settings_numpy = (settings or {}).copy() + if settings is not None and "pareto_probs" in settings: + settings_numpy["pareto_probs"] = { + participants_codes.columns.tolist().index(k): v + for k, v in settings_numpy["pareto_probs"].items() + } # Run allocation with numpy arrays. result: AllocatorResult = allocate_numpy( @@ -58,7 +71,7 @@ def allocate_pandas(participants: pd.DataFrame, manuals=manuals, algorithm=algorithm, progress_func=progress_func, - settings=settings, + settings=settings_numpy, ) # Generate pandas dataframe based on numpy results and return. @@ -66,19 +79,19 @@ def allocate_pandas(participants: pd.DataFrame, [ # This is just to ensure that the two columns allocation # and group will be the first columns of the dataframe. - pd.DataFrame(columns=['allocation', 'group']) - ] + [ - participants - .loc[participants.index[group]] - .reset_index( - names=[f"participant {l}" for l in participants.index.names] - if participants.index.nlevels > 1 else - 'participant' - ) - .assign( - allocation=al_id, - group=gr_id, - ) + pd.DataFrame(columns=["allocation", "group"]) + ] + + [ + participants.loc[participants.index[group]] + .reset_index( + names=[f"participant {l}" for l in participants.index.names] + if participants.index.nlevels > 1 + else "participant" + ) + .assign( + allocation=al_id, + group=gr_id, + ) for al_id, allocation in enumerate(result.ensemble) for gr_id, group in enumerate(allocation) ], @@ -88,11 +101,13 @@ def allocate_pandas(participants: pd.DataFrame, return ret_part ret_groups = pd.DataFrame.from_records( [ - pd.Series({ - 'allocation': al_id, - 'group': gr_id, - 'participant': participants.index[group].tolist(), - }) + pd.Series( + { + "allocation": al_id, + "group": gr_id, + "participant": participants.index[group].tolist(), + } + ) for al_id, allocation in enumerate(result.ensemble) for gr_id, group in enumerate(allocation) ] @@ -100,19 +115,25 @@ def allocate_pandas(participants: pd.DataFrame, return ret_part, ret_groups, result -@pd.api.extensions.register_dataframe_accessor('groupselect') +@pd.api.extensions.register_dataframe_accessor("groupselect") class GroupSelectAccessor: + """Pandas ``.groupselect`` accessor for allocating a DataFrame.""" + def __init__(self, df: pd.DataFrame): + """Bind the accessor to the DataFrame it was accessed on.""" self._df = df - def allocate(self, - fields: dict[Hashable, FieldMode], - n_part_per_group: int | Iterable[int], - manuals: None | dict[int, int] = None, - algorithm: Algorithm | str = Algorithm.Legacy, - progress_func: None | Callable = None, - settings: None | dict = None, - return_full: bool = False) -> pd.DataFrame: + def allocate( + self, + fields: dict[Hashable, FieldMode], + n_part_per_group: int | Iterable[int], + manuals: None | dict[int, int] = None, + algorithm: Algorithm | str = Algorithm.Legacy, + progress_func: None | Callable = None, + settings: None | dict = None, + return_full: bool = False, + ) -> pd.DataFrame: + """Allocate the bound DataFrame's participants into groups.""" return allocate_pandas( participants=self._df, fields=fields, diff --git a/groupselect/allocation.py b/groupselect/allocation.py index f91bc02..01569d2 100644 --- a/groupselect/allocation.py +++ b/groupselect/allocation.py @@ -1,25 +1,108 @@ +"""Data structures for allocations and metrics computed over them.""" + +from __future__ import annotations + +from collections import Counter +from itertools import combinations +from math import comb + +import numpy as np + +try: + import pandas as pd +except ImportError: + pd = None + + class ParticipantGroup(list[int]): + """A single group of participant IDs.""" + pass class Allocation(list[ParticipantGroup]): + """A full allocation of participants into groups.""" + pass class AllocationEnsemble(list[Allocation]): + """A collection of allocations, e.g. produced over multiple rounds.""" + + def get_pids(self): + """Return the set of all participant IDs across all allocations.""" + return { + p_id + for allocation in self + for group in allocation + for p_id in group + if p_id == p_id + } + + def calc_total_number_pairs(self): + """Return the number of possible participant pairs.""" + pids = self.get_pids() + return comb(len(pids), 2) + + def calc_pair_counts(self) -> Counter: + """Count how often each pair of participants shares a group.""" + pair_counts = Counter() + for allocation in self: + for group in allocation: + for pair in combinations(group, 2): + pair_counts[frozenset(pair)] += 1 + return pair_counts + + def calc_pair_occurrences(self) -> Counter: + """Count how many pairs meet 0, 1, 2, ... times across allocations.""" + # Calculate occurrences from pair counter. + pair_counts = self.calc_pair_counts() + occurrences = Counter(pair_counts.values()) + + # Add entry on non-observed pairs. + total_pairs = self.calc_total_number_pairs() + observed_pairs = len(pair_counts) + occurrences[0] = total_pairs - observed_pairs + + return occurrences + + def calc_meeting_rel_score(self) -> float: + """Return the share of participant pairs that meet at least once.""" + occurrences = self.calc_pair_occurrences() + return 1 - occurrences[0] / sum(occurrences.values()) + + def calc_meeting_norm_score(self) -> float: + """Return the meeting score normalised to its achievable range.""" + occurrences = self.calc_pair_occurrences() + rel_score = sum(v for k, v in occurrences.items() if k) + + number_pairs_per_allocation = [ + sum(comb(len(group), 2) for group in allocation) + for allocation in self + ] + min_score = max(number_pairs_per_allocation) + max_score = sum(number_pairs_per_allocation) + + return ( + ((rel_score - min_score) / (max_score - min_score)) + if max_score > min_score + else 1.0 + ) + def calc_n_meetings_alo(self) -> int: + """Return the total number of distinct meeting partners, summed.""" return sum( - len(p_stats) - for p_id, p_stats in self.calc_meetings().items() + len(p_stats) for p_id, p_stats in self.calc_meetings().items() ) def calc_meetings(self) -> dict[int, dict[int, int]]: + """Return how often each participant met every other one.""" p_ids = { p_id for allocation in self for group in allocation for p_id in group - if p_id + if p_id == p_id } meetings = {} @@ -36,7 +119,50 @@ def calc_meetings(self) -> dict[int, dict[int, int]]: return meetings + def _calc_diversity_score_single_field(self, values_field: np.array): + values, inverse = np.unique(values_field, return_inverse=True) + + counts_full = np.bincount(inverse) + share_full = counts_full / len(values_field) + + ret = 0.0 + for allocation in self: + for group in allocation: + group_codes = inverse[group] + + counts_group = np.bincount( + group_codes, + minlength=len(values), + ) + + share_group = counts_group / len(group) + + ret += np.abs(share_full - share_group).sum() + ret /= len(allocation) + + return ret / len(self) + + def calc_diversity_score( + self, participants_data: np.ndarray[int] | pd.DataFrame + ) -> float: + """Return the summed diversity score across all fields.""" + if pd is not None and isinstance(participants_data, pd.DataFrame): + participants_data = ( + participants_data.astype("category") + .apply(lambda col: col.cat.codes) + .to_numpy() + ) + + # Sum over fields. + return sum( + self._calc_diversity_score_single_field(col) + for col in participants_data.T + ) + class AllocatorResult: + """Result of an allocator run, wrapping the resulting ensemble.""" + def __init__(self, ensemble: None | AllocationEnsemble = None): + """Store the allocation ensemble produced by the allocator.""" self.ensemble = ensemble diff --git a/groupselect/examples/__init__.py b/groupselect/examples/__init__.py index 5521401..bb90640 100644 --- a/groupselect/examples/__init__.py +++ b/groupselect/examples/__init__.py @@ -1,29 +1,32 @@ +"""Example participant datasets loaded from CSV files for demos and tests.""" + from pathlib import Path try: import pandas as pd + HAS_PANDAS: bool = True except ImportError: pd = None HAS_PANDAS: bool = False -DATA_DIR = Path(__file__).parent / 'data' +DATA_DIR = Path(__file__).parent / "data" if HAS_PANDAS: example_data_np = {} example_data_pd = {} - for file in DATA_DIR.glob('*.csv'): + for file in DATA_DIR.glob("*.csv"): basename = file.stem example_data_pd[basename] = pd.read_csv(file) example_data_pd[basename].set_index( - 'ID' if 'ID' in example_data_pd[basename] else 'name', + "ID" if "ID" in example_data_pd[basename] else "name", inplace=True, ) example_data_np[basename] = ( example_data_pd[basename] - .astype('category') + .astype("category") .apply(lambda col: col.cat.codes) .to_numpy() ) diff --git a/groupselect/field_mode.py b/groupselect/field_mode.py index afde5ec..aded307 100644 --- a/groupselect/field_mode.py +++ b/groupselect/field_mode.py @@ -1,8 +1,12 @@ +"""Modes controlling how a participant field is treated during allocation.""" + from enum import Enum, unique @unique class FieldMode(Enum): + """How a participant field should be used when forming groups.""" + Ignore = 0 Diversify = 1 Cluster = 2 diff --git a/groupselect/tests/test_pandas.py b/groupselect/tests/test_pandas.py deleted file mode 100644 index fe2024e..0000000 --- a/groupselect/tests/test_pandas.py +++ /dev/null @@ -1,38 +0,0 @@ -import pytest - -import pandas as pd - -from groupselect.examples import example_data_pd - - -# def test_100i_4j_3k_3d_data(): -# df = example_data_pd['100i_4j_3k_3d_data'] -# for fields in ({'age': 'diversify', 'gender': 'diversify'}, -# {'gender': 'diversify', 'age': 'diversify'}, -# {'age': 'diversify', 'gender': 'diversify', -# 'region': 'diversify', 'nation': 'diversify'}): -# for n_part_per_group in (3 * [6], 4 * [8],): -# res = df.groupselect.allocate( -# fields=fields, -# n_part_per_group=n_part_per_group, -# ) -# with pd.option_context('display.max_rows', None, 'display.max_columns', None): -# print(res) - - -def test_philipps_example_data(): - df = example_data_pd['philipps_example_data'] - for fields in (#{'age': 'diversify', 'gender': 'diversify'}, - #{'gender': 'diversify', 'age': 'diversify'}, - #{'age': 'diversify', 'gender': 'diversify', - # 'region': 'diversify', 'ethnicity': 'diversify', - # 'education': 'diversify'}, - {'age': 'diversify', 'gender': 'diversify', - 'photo consent': 'cluster'},): - for n_part_per_group in (3 * [6], 4 * [8],): - res = df.groupselect.allocate( - fields=fields, - n_part_per_group=n_part_per_group, - ) - with pd.option_context('display.max_rows', None, 'display.max_columns', None): - print(res) diff --git a/pyproject.toml b/pyproject.toml index dd9a6cb..8499575 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,7 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + [project] name = "groupselect" version = "0.1.0" @@ -12,7 +16,21 @@ dependencies = [ "numpy>=2.3.4", ] +[project.optional-dependencies] +pandas = ["pandas>=2.0"] + [dependency-groups] dev = [ - "pytest>=9.0.2", + "pytest>=9.0.0", + "ruff>=0.16.5", ] + +[tool.ruff] +line-length = 79 + +[tool.ruff.lint] +select = ["D"] +extend-select = ["E501"] + +[tool.ruff.lint.pydocstyle] +convention = "pep257" diff --git a/groupselect/tests/test_load_examples.py b/tests/test_load_examples.py similarity index 52% rename from groupselect/tests/test_load_examples.py rename to tests/test_load_examples.py index 5a56253..9be41f6 100644 --- a/groupselect/tests/test_load_examples.py +++ b/tests/test_load_examples.py @@ -1,7 +1,10 @@ +"""Tests for loading the example datasets.""" + import pytest def test_load_examples(): + """Check that the examples package can be imported without error.""" try: import groupselect.examples except Exception as ex: diff --git a/tests/test_pandas.py b/tests/test_pandas.py new file mode 100644 index 0000000..6694498 --- /dev/null +++ b/tests/test_pandas.py @@ -0,0 +1,31 @@ +"""Tests for the pandas ``.groupselect`` accessor.""" + +import pytest + +import pandas as pd + +from groupselect.examples import example_data_pd + + +def test_philipps_example_data(): + """Allocate Philipp's example dataset and print the result.""" + df = example_data_pd["philipps_example_data"] + for fields in ( + { + "age": "diversify", + "gender": "diversify", + "photo consent": "cluster", + }, + ): + for n_part_per_group in ( + 3 * [6], + 4 * [8], + ): + res = df.groupselect.allocate( + fields=fields, + n_part_per_group=n_part_per_group, + ) + with pd.option_context( + "display.max_rows", None, "display.max_columns", None + ): + print(res)