Update genetic_algorithm_optimization.py

genetic_algorithm/genetic_algorithm_optimization.py

@@ -1,7 +1,6 @@
 import random
 from collections.abc import Callable, Sequence
 from concurrent.futures import ThreadPoolExecutor
-
 import numpy as np
 
 # Parameters
@@ -40,7 +39,25 @@ class GeneticAlgorithm:
         self.population = self.initialize_population()
 
     def initialize_population(self) -> list[np.ndarray]:
-        """Initialize the population with random individuals within the search space."""
+        """
+        Initialize the population with random individuals within the search space.
+
+        Example:
+        >>> ga = GeneticAlgorithm(
+        ...     function=lambda x, y: x**2 + y**2,
+        ...     bounds=[(-10, 10), (-10, 10)],
+        ...     population_size=5,
+        ...     generations=10,
+        ...     mutation_prob=0.1,
+        ...     crossover_rate=0.8,
+        ...     maximize=False
+        ... )
+        >>> len(ga.initialize_population())
+        5  # The population size should be equal to 5.
+        >>> all(len(ind) == 2 for ind in ga.initialize_population())  
+        # Each individual should have 2 variables
+        True
+        """
         return [
             rng.uniform(
                 low=[self.bounds[j][0] for j in range(self.dim)],
@@ -50,14 +67,58 @@ class GeneticAlgorithm:
         ]
 
     def fitness(self, individual: np.ndarray) -> float:
-        """Calculate the fitness value (function value) for an individual."""
+        """
+        Calculate the fitness value (function value) for an individual.
+
+        Example:
+        >>> ga = GeneticAlgorithm(
+        ...     function=lambda x, y: x**2 + y**2,
+        ...     bounds=[(-10, 10), (-10, 10)],
+        ...     population_size=10,
+        ...     generations=10,
+        ...     mutation_prob=0.1,
+        ...     crossover_rate=0.8,
+        ...     maximize=False
+        ... )
+        >>> individual = np.array([1.0, 2.0])
+        >>> ga.fitness(individual)
+        5.0  # The fitness should be 1^2 + 2^2 = 5
+        >>> ga.maximize = True
+        >>> ga.fitness(individual)
+        -5.0  # The fitness should be -5 when maximizing
+        """
         value = float(self.function(*individual))  # Ensure fitness is a float
         return value if self.maximize else -value  # If minimizing, invert the fitness
 
     def select_parents(
         self, population_score: list[tuple[np.ndarray, float]]
     ) -> list[np.ndarray]:
-        """Select top N_SELECTED parents based on fitness."""
+        """
+        Select top N_SELECTED parents based on fitness.
+
+        Example:
+        >>> ga = GeneticAlgorithm(
+        ...     function=lambda x, y: x**2 + y**2,
+        ...     bounds=[(-10, 10), (-10, 10)],
+        ...     population_size=10,
+        ...     generations=10,
+        ...     mutation_prob=0.1,
+        ...     crossover_rate=0.8,
+        ...     maximize=False
+        ... )
+        >>> population_score = [
+        ...     (np.array([1.0, 2.0]), 5.0),
+        ...     (np.array([-1.0, -2.0]), 5.0),
+        ...     (np.array([0.0, 0.0]), 0.0),
+        ... ]
+        >>> selected_parents = ga.select_parents(population_score)
+        >>> len(selected_parents)
+        2  # Should select the two parents with the best fitness scores.
+        >>> np.array_equal(selected_parents[0], np.array([1.0, 2.0]))  # Parent 1 should be [1.0, 2.0]
+        True
+        >>> np.array_equal(selected_parents[1], np.array([-1.0, -2.0]))  # Parent 2 should be [-1.0, -2.0]
+        True
+        """
         population_score.sort(key=lambda score_tuple: score_tuple[1], reverse=True)
         selected_count = min(N_SELECTED, len(population_score))
         return [ind for ind, _ in population_score[:selected_count]]
@@ -67,11 +128,13 @@ class GeneticAlgorithm:
     ) -> tuple[np.ndarray, np.ndarray]:
         """
         Perform uniform crossover between two parents to generate offspring.
+
         Args:
             parent1 (np.ndarray): The first parent.
             parent2 (np.ndarray): The second parent.
         Returns:
             tuple[np.ndarray, np.ndarray]: The two offspring generated by crossover.
+        
         Example:
         >>> ga = GeneticAlgorithm(
         ...     lambda x, y: -(x**2 + y**2),
@@ -92,10 +155,13 @@ class GeneticAlgorithm:
     def mutate(self, individual: np.ndarray) -> np.ndarray:
         """
         Apply mutation to an individual.
+
         Args:
             individual (np.ndarray): The individual to mutate.
+
         Returns:
             np.ndarray: The mutated individual.
+
         Example:
         >>> ga = GeneticAlgorithm(
         ...     lambda x, y: -(x**2 + y**2),
@@ -115,9 +181,11 @@ class GeneticAlgorithm:
     def evaluate_population(self) -> list[tuple[np.ndarray, float]]:
         """
         Evaluate the fitness of the entire population in parallel.
+
         Returns:
             list[tuple[np.ndarray, float]]:
                 The population with their respective fitness values.
+
         Example:
         >>> ga = GeneticAlgorithm(
         ...     lambda x, y: -(x**2 + y**2),
@@ -141,11 +209,33 @@ class GeneticAlgorithm:
                 )
             )
 
-    def evolve(self, verbose=True) -> np.ndarray:
+    def evolve(self, verbose: bool = True) -> np.ndarray:
         """
         Evolve the population over the generations to find the best solution.
+
+        Args:
+            verbose (bool): If True, prints the progress of the generations.
+
         Returns:
             np.ndarray: The best individual found during the evolution process.
+
+        Example:
+        >>> ga = GeneticAlgorithm(
+        ...     function=lambda x, y: x**2 + y**2,
+        ...     bounds=[(-10, 10), (-10, 10)],
+        ...     population_size=10,
+        ...     generations=10,
+        ...     mutation_prob=0.1,
+        ...     crossover_rate=0.8,
+        ...     maximize=False
+        ... )
+        >>> best_solution = ga.evolve(verbose=False)
+        >>> len(best_solution)
+        2  # The best solution should be a 2-element array (var_x, var_y)
+        >>> isinstance(best_solution[0], float)  # First element should be a float
+        True
+        >>> isinstance(best_solution[1], float)  # Second element should be a float
+        True
         """
         for generation in range(self.generations):
             # Evaluate population fitness (multithreaded)
@@ -186,6 +276,7 @@ def target_function(var_x: float, var_y: float) -> float:
         var_y (float): The y-coordinate.
     Returns:
         float: The value of the function at (var_x, var_y).
+    
     Example:
     >>> target_function(0, 0)
     0