рисование дерева

lab4/gp/ga.py

@@ -6,11 +6,12 @@ from copy import deepcopy
 from dataclasses import asdict, dataclass
 from typing import Callable
 
+import graphviz
 import numpy as np
 from matplotlib import pyplot as plt
-from numpy.typing import NDArray
 
 from .chromosome import Chromosome
+from .node import Node
 from .types import Fitnesses, Population
 
 type FitnessFn = Callable[[Chromosome], float]
@@ -66,6 +67,7 @@ class Generation:
     number: int
     best: Chromosome
     best_fitness: float
+    avg_fitness: float
     population: Population
     fitnesses: Fitnesses
 
@@ -148,27 +150,45 @@ def eval_population(population: Population, fitness_func: FitnessFn) -> Fitnesse
     return np.array([fitness_func(chrom) for chrom in population])
 
 
+def render_tree_to_graphviz(
+    node: Node, graph: graphviz.Digraph, node_id: str = "0"
+) -> None:
+    """Рекурсивно добавляет узлы дерева в graphviz граф."""
+    graph.node(node_id, label=node.value.name)
+
+    for i, child in enumerate(node.children):
+        child_id = f"{node_id}_{i}"
+        render_tree_to_graphviz(child, graph, child_id)
+        graph.edge(node_id, child_id)
+
+
 def save_generation(
     generation: Generation, history: list[Generation], config: GARunConfig
 ) -> None:
+    """Сохраняет визуализацию лучшей хромосомы поколения в виде дерева."""
     os.makedirs(config.results_dir, exist_ok=True)
 
-    fig = plt.figure(figsize=(7, 7))
-    fig.suptitle(
-        f"Поколение #{generation.number}. "
-        f"Лучшая особь: {generation.best_fitness:.0f}. "
-        f"Среднее значение: {np.mean(generation.fitnesses):.0f}",
-        fontsize=14,
-        y=0.95,
+    # Создаем граф для визуализации дерева
+    dot = graphviz.Digraph(comment=f"Generation {generation.number}")
+    dot.attr(rankdir="TB")  # Top to Bottom direction
+    dot.attr("node", shape="circle", style="filled", fillcolor="lightblue")
+
+    # Добавляем заголовок
+    depth = generation.best.root.get_depth()
+    title = (
+        f"Поколение #{generation.number}\\n"
+        f"Лучшая особь: {generation.best_fitness:.4f}\\n"
+        f"Глубина дерева: {depth}"
     )
+    dot.attr(label=title, labelloc="t", fontsize="14")
 
-    # Рисуем
-    ...
+    # Рендерим дерево
+    render_tree_to_graphviz(generation.best.root, dot)
 
-    filename = f"generation_{generation.number:03d}.png"
-    path_png = os.path.join(config.results_dir, filename)
-    fig.savefig(path_png, dpi=150, bbox_inches="tight")
-    plt.close(fig)
+    # Сохраняем
+    filename = f"generation_{generation.number:03d}"
+    filepath = os.path.join(config.results_dir, filename)
+    dot.render(filepath, format="png", cleanup=True)
 
 
 def genetic_algorithm(config: GARunConfig) -> GARunResult:
@@ -216,6 +236,7 @@ def genetic_algorithm(config: GARunConfig) -> GARunResult:
             number=generation_number,
             best=population[best_index],
             best_fitness=fitnesses[best_index],
+            avg_fitness=float(np.mean(fitnesses)),
             # population=deepcopy(population),
             population=[],
             # fitnesses=deepcopy(fitnesses),
@@ -301,41 +322,61 @@ def genetic_algorithm(config: GARunConfig) -> GARunResult:
     end = time.perf_counter()
 
     assert best is not None, "Best was never set"
-    return GARunResult(
+    result = GARunResult(
         len(history),
         best,
         history,
         (end - start) * 1000.0,
     )
 
+    # Автоматически строим графики истории фитнеса
+    if config.save_generations:
+        plot_fitness_history(result, save_dir=config.results_dir)
 
-def plot_fitness_history(result: GARunResult, save_path: str | None = None) -> None:
-    """Рисует график изменения лучших и средних значений фитнеса по поколениям."""
+    return result
+
+
+def plot_fitness_history(result: GARunResult, save_dir: str | None = None) -> None:
+    """Рисует графики изменения лучших и средних значений фитнеса по поколениям.
+
+    Создает два отдельных графика:
+    - fitness_best.png - график лучших значений
+    - fitness_avg.png - график средних значений
+    """
     generations = [gen.number for gen in result.history]
     best_fitnesses = [gen.best_fitness for gen in result.history]
-    avg_fitnesses = [np.mean(gen.fitnesses) for gen in result.history]
+    avg_fitnesses = [gen.avg_fitness for gen in result.history]
 
-    fig, ax = plt.subplots(figsize=(10, 6))
+    # График лучших значений
+    fig_best, ax_best = plt.subplots(figsize=(10, 6))
+    ax_best.plot(generations, best_fitnesses, linewidth=2, color="blue")
+    ax_best.set_xlabel("Поколение", fontsize=12)
+    ax_best.set_ylabel("Лучшее значение фитнес-функции", fontsize=12)
+    ax_best.set_title("Лучшее значение фитнеса по поколениям", fontsize=14)
+    ax_best.grid(True, alpha=0.3)
 
-    ax.plot(
-        generations, best_fitnesses, label="Лучшее значение", linewidth=2, color="blue"
-    )
-    ax.plot(
-        generations,
-        avg_fitnesses,
-        label="Среднее значение",
-        linewidth=2,
-        color="orange",
-    )
-
-    ax.set_xlabel("Поколение", fontsize=12)
-    ax.set_ylabel("Значение фитнес-функции", fontsize=12)
-    ax.legend(fontsize=11)
-    ax.grid(True, alpha=0.3)
-
-    if save_path:
-        fig.savefig(save_path, dpi=150, bbox_inches="tight")
-        print(f"График сохранен в {save_path}")
+    if save_dir:
+        best_path = os.path.join(save_dir, "fitness_best.png")
+        fig_best.savefig(best_path, dpi=150, bbox_inches="tight")
+        print(f"График лучших значений сохранен в {best_path}")
     else:
         plt.show()
-        plt.close(fig)
+
+    plt.close(fig_best)
+
+    # График средних значений
+    fig_avg, ax_avg = plt.subplots(figsize=(10, 6))
+    ax_avg.plot(generations, avg_fitnesses, linewidth=2, color="orange")
+    ax_avg.set_xlabel("Поколение", fontsize=12)
+    ax_avg.set_ylabel("Среднее значение фитнес-функции", fontsize=12)
+    ax_avg.set_title("Среднее значение фитнеса по поколениям", fontsize=14)
+    ax_avg.grid(True, alpha=0.3)
+
+    if save_dir:
+        avg_path = os.path.join(save_dir, "fitness_avg.png")
+        fig_avg.savefig(avg_path, dpi=150, bbox_inches="tight")
+        print(f"График средних значений сохранен в {avg_path}")
+    else:
+        plt.show()
+
+    plt.close(fig_avg)