diff --git a/pyproject.toml b/pyproject.toml index 1bb3e82..5085f50 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ keywords = ["game theory", "tikz", "visualization", "trees", "economics"] # Required runtime dependencies (previously optional under the 'jupyter' extra) -dependencies = ["jupyter-tikz", "ipykernel"] +dependencies = ["jupyter-tikz", "ipykernel", "distinctipy"] [project.optional-dependencies] dev = ["pytest>=7.0.0", "pytest-cov", "nbformat", "nbclient", "ipykernel", "pygambit"] diff --git a/src/draw_tree/core.py b/src/draw_tree/core.py index 779945c..58fbe13 100644 --- a/src/draw_tree/core.py +++ b/src/draw_tree/core.py @@ -11,6 +11,7 @@ import subprocess import tempfile import re +import distinctipy from typing import TYPE_CHECKING from numpy import save @@ -79,14 +80,19 @@ def get_player_color(player: int, color_scheme: str = "default") -> str: Get the TeX color macro name for a given player number. Args: - player: Player number (1-6 for regular players). - color_scheme: Optional color scheme name. + player: Player number (0 for chance, 1-6 for regular players with + "gambit" scheme, or any positive integer for "distinctipy" and + "colorblind" schemes). + color_scheme: Color scheme name. One of "default", "gambit", + "distinctipy", or "colorblind". Returns: TeX color macro name for the player, or "black" as fallback. + + Raises: + ValueError: If the "gambit" scheme is used with more than 6 players. """ if color_scheme == "gambit": - # Color mapping for up to 6 players color_map = { 0: "\\chancecolor", 1: "\\playeronecolor", @@ -96,24 +102,84 @@ def get_player_color(player: int, color_scheme: str = "default") -> str: 5: "\\playerfivecolor", 6: "\\playersixcolor", } + if player < 0: + return "black" # no player assigned yet + if player not in color_map: + raise ValueError( + f"The 'gambit' color scheme only supports up to 6 players " + f"(got player {player}). Consider using the 'distinctipy' " + f"or 'colorblind' color scheme for games with more players." + ) + return color_map[player] + + elif color_scheme in ("distinctipy", "colorblind"): + if player == 0: + return "\\chancecolor" + elif player > 0: + return f"\\player{player}color" - return color_map.get(player, "black") return "black" -def color_definitions() -> list[str]: - return [ +def color_definitions(color_scheme: str = "default", num_players: int = 6) -> list[str]: + """ + Generate LaTeX color macro definitions for game tree players. + + Produces ``\\definecolor`` and ``\\newcommand`` lines that are injected + into the TikZ preamble so that player-color macros (e.g. + ``\\playeronecolor``, ``\\player7color``) resolve correctly. + + Args: + color_scheme: One of "default", "gambit", "distinctipy", or + "colorblind". + num_players: Number of players that need colours. Ignored for + "default" and "gambit" (which have fixed palettes). + + Returns: + List of LaTeX definition strings. + """ + # Chance color is shared across all schemes + defs = [ "\\definecolor{chancecolorrgb}{RGB}{117,145,56}", - "\\definecolor{gambitredrgb}{RGB}{234,51,35}", "\\newcommand\\chancecolor{chancecolorrgb}", - "\\newcommand\\playeronecolor{gambitredrgb}", - "\\newcommand\\playertwocolor{blue}", - "\\newcommand\\playerthreecolor{orange}", - "\\newcommand\\playerfourcolor{purple}", - "\\newcommand\\playerfivecolor{cyan}", - "\\newcommand\\playersixcolor{magenta}", ] + if color_scheme == "gambit": + defs.extend([ + "\\definecolor{gambitredrgb}{RGB}{234,51,35}", + "\\newcommand\\playeronecolor{gambitredrgb}", + "\\newcommand\\playertwocolor{blue}", + "\\newcommand\\playerthreecolor{orange}", + "\\newcommand\\playerfourcolor{purple}", + "\\newcommand\\playerfivecolor{cyan}", + "\\newcommand\\playersixcolor{magenta}", + ]) + + elif color_scheme in ("distinctipy", "colorblind"): + # Chance color in 0-1 float format for exclusion + chance_rgb = (117 / 255, 145 / 255, 56 / 255) + try: + colorblind_type = ( + "Deuteranomaly" if color_scheme == "colorblind" else None + ) + colors = distinctipy.get_colors( + num_players, + exclude_colors=[(0, 0, 0), (1, 1, 1), chance_rgb], + rng=42, + colorblind_type=colorblind_type, + ) + for i, color in enumerate(colors): + r, g, b = [int(c * 255) for c in color] + p_num = i + 1 + defs.extend([ + f"\\definecolor{{p{p_num}rgb}}{{RGB}}{{{r},{g},{b}}}", + f"\\newcommand\\player{p_num}color{{p{p_num}rgb}}", + ]) + except Exception as e: + print(f"Warning: Failed to generate {color_scheme} colors: {e}") + + return defs + def outall(stream: Optional[List[str]] = None) -> None: """ Output stream to stdout. @@ -1573,6 +1639,28 @@ def generate_tikz( shared_terminal_depth=shared_terminal_depth, ) + # Determine the number of players for dynamic color schemes + num_players = 0 + if not isinstance(game, str): + try: + num_players = len(game.players) + except AttributeError: + num_players = 6 + else: + try: + player_nums = set() + for line in readfile(ef_file): + if line.startswith("player"): + try: + p = int(line.split()[1]) + if p > 0: + player_nums.add(p) + except (IndexError, ValueError): + pass + num_players = len(player_nums) if player_nums else 6 + except Exception: + num_players = 6 + # Step 1: Generate the tikzpicture content using ef_to_tex logic tikz_picture_content = ef_to_tex(ef_file, scale_factor, show_grid, color_scheme, action_label_position) @@ -1596,7 +1684,7 @@ def generate_tikz( f"\\treethickn{edge_thickness}pt", ] # Step 2a: Define player color macros - macro_definitions.extend(color_definitions()) + macro_definitions.extend(color_definitions(color_scheme, num_players)) # Step 3: Combine everything into complete TikZ code tikz_code = """% TikZ code with built-in styling for game trees