Repository: OskarLiew/PyGenOpt

Issue:
Add unit tests for functions
Needs tests for:

* crossover
* evaluation
* mutation
* population (not prio)
* selection

Discussion:

Patch:
diff --git a/.gitignore b/.gitignore
index b6e4761..e02647d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# vscode settings
+.vscode/
\ No newline at end of file
diff --git a/README.md b/README.md
index 32889f6..7af9e51 100644
--- a/README.md
+++ b/README.md
@@ -19,7 +19,9 @@ Example code to minimize [BIC](https://en.wikipedia.org/wiki/Bayesian_informatio
 ```python
 from functools import partial
 import statsmodels.api as sm
+from sklearn import datasets
 import numpy as np
+import pandas as pd
 from genopt import GeneticOptimizer
 
 
@@ -41,12 +43,10 @@ def linear_regression_minimize_bic(chromosome, X, y):
     return -results.bic
 
 
-# Load data. There are not very many variables in this dataset, so there
-# is a large chance that the best subset is found in the first generation
-data = sm.datasets.longley.load_pandas()
-X = sm.add_constant(data.exog)
-y = data.endog
-feature_names = X.columns
+# Load Boston housing dataset
+data = datasets.load_boston()
+X = pd.DataFrame(data=data["data"], columns=data["feature_names"])
+y = data["target"]
 
 # Setup GeneticOptimizer
 go = GeneticOptimizer(
@@ -66,4 +66,5 @@ X_subset = get_X_subset(best_chromosome, X)
 lr = sm.OLS(y, X_subset)
 results = lr.fit()
 print(results.summary())
+
 ```
diff --git a/examples/linear_regression_variable_selection.py b/examples/linear_regression_variable_selection.py
index f45e929..f67b492 100644
--- a/examples/linear_regression_variable_selection.py
+++ b/examples/linear_regression_variable_selection.py
@@ -1,49 +1,55 @@
 from functools import partial
 import statsmodels.api as sm
+from sklearn import datasets
 import numpy as np
+import pandas as pd
 from genopt import GeneticOptimizer
 
 
 # Helper function to get the columns defined by the chromosome
-def get_X_subset(chromosome, X):
+def get_x_subset(chromosome, x_data):
     mask = chromosome.astype(bool)
-    return X.iloc[:, mask]
+    return x_data.iloc[:, mask]
 
 
 # Objective function to optimize
-def linear_regression_minimize_bic(chromosome, X, y):
+def linear_regression_minimize_bic(chromosome, x_data, targets):
     if np.all(chromosome == 0):
         return -1e9
 
-    X_subset = get_X_subset(chromosome, X)
-    lr = sm.OLS(y, X_subset)
-    results = lr.fit()
+    x_subset = get_x_subset(chromosome, x_data)
+    regression_model = sm.OLS(targets, x_subset)
+    results = regression_model.fit()
 
     return -results.bic
 
 
-# Load data. There are not very many variables in this dataset, so there
-# is a large chanse that the best subset is found in the first generation
-data = sm.datasets.longley.load_pandas()
-X = sm.add_constant(data.exog)
-y = data.endog
-feature_names = X.columns
-
-# Setup GeneticOptimizer
-go = GeneticOptimizer(
-    n_vars=X.shape[1],
-    popsize=100,
-    n_gen=10,
-    objective_function=partial(linear_regression_minimize_bic, X=X, y=y),
-    encoding="discrete",
-    var_range=(0, 1),
-    var_size=1,
-)
-
-best_chromosome = go.optimize()
-
-# Show results of optimization
-X_subset = get_X_subset(best_chromosome, X)
-lr = sm.OLS(y, X_subset)
-results = lr.fit()
-print(results.summary())
+def main():
+    # Load Boston housing dataset
+    data = datasets.load_boston()
+    x_data = pd.DataFrame(data=data["data"], columns=data["feature_names"])
+    targets = data["target"]
+
+    # Setup GeneticOptimizer
+    optimizer = GeneticOptimizer(
+        n_vars=x_data.shape[1],
+        popsize=100,
+        objective_function=partial(
+            linear_regression_minimize_bic, x_data=x_data, targets=targets
+        ),
+        encoding="discrete",
+        var_range=(0, 1),
+        var_size=1,
+    )
+
+    best_chromosome = optimizer.optimize(10)
+
+    # Show results of optimization
+    x_subset = get_x_subset(best_chromosome, x_data)
+    regression_model = sm.OLS(targets, x_subset)
+    results = regression_model.fit()
+    print(results.summary())
+
+
+if __name__ == "__main__":
+    main()
diff --git a/genopt/crossover.py b/genopt/crossover.py
index ca41d59..256b038 100644
--- a/genopt/crossover.py
+++ b/genopt/crossover.py
@@ -1,14 +1,18 @@
 import numpy as np
 
 
-def one_way_crossover(population, n_vars, popsize):
-    for j in np.arange(popsize, step=2):
-        group1 = population[j, :]
-        group2 = population[j + 1, :]
-        cross_point = np.random.randint(n_vars)
-        group1[:cross_point] = group2[:cross_point]
-        group2[cross_point:] = group1[cross_point:]
-        population[j, :] = group1
-        population[j + 1, :] = group2
-
-    return population
+def one_way_crossover(population: np.ndarray) -> np.ndarray:
+    new_population = np.atleast_2d(population.copy())
+    popsize, chromosome_length = new_population.shape
+    for j in np.arange(popsize - 1, step=2):
+
+        # Take two chromosomes and find a crossover point
+        chromosome1 = new_population[j, :].copy()
+        chromosome2 = new_population[j + 1, :].copy()
+        cross_point = np.random.randint(1, chromosome_length - 1)
+
+        # Swap the chromosomes at the crossover point and update the population
+        new_population[j, :cross_point] = chromosome2[:cross_point]
+        new_population[j + 1, :cross_point] = chromosome1[:cross_point]
+
+    return new_population
diff --git a/genopt/evaluation.py b/genopt/evaluation.py
index 472a56b..c6467b8 100644
--- a/genopt/evaluation.py
+++ b/genopt/evaluation.py
@@ -1,19 +1,21 @@
+from typing import Callable, Tuple
 import numpy as np
 
 
-def decode(population, n_vars, encoding, var_range, var_size):
+def decode_discrete(
+    population: np.ndarray, n_vars: int, var_range: Tuple[float], var_size: int
+) -> float:
     # Special case where there is no need to decode
-    if (
-        var_range[0] == 0 and var_range[1] == 1 and var_size == 1
-    ) or encoding == "real":
+    if var_range[0] == 0 and var_range[1] == 1 and var_size == 1:
         return population
 
-    population = population.reshape((-1, n_vars * var_size))
+    population = np.atleast_2d(population.copy())
     popsize = population.shape[0]
     decoded = np.zeros((popsize, n_vars))
     var_min = var_range[0]
     var_max = var_range[1]
 
+    # Has room for optimization
     for i_indiv in range(popsize):
         for i_var in range(n_vars):
             var_bits = population[
@@ -29,13 +31,8 @@ def decode(population, n_vars, encoding, var_range, var_size):
     return decoded
 
 
-def evaluate_pop(population, n_vars, objective_function, var_range, var_size, encoding):
-    if encoding == "discrete":
-        variables = decode(population, n_vars, encoding, var_range, var_size)
-    else:
-        variables = population
-
-    popsize = population.shape[0]
+def evaluate(variables: np.ndarray, objective_function: Callable) -> np.ndarray:
+    variables = np.atleast_2d(variables.copy())
+    popsize = variables.shape[0]
     fitness = [objective_function(variables[i, :]) for i in range(popsize)]
-    i_max = np.argmax(fitness)
-    return np.array(fitness), i_max
+    return np.array(fitness)
diff --git a/genopt/go.py b/genopt/go.py
index afec555..daf9ed8 100644
--- a/genopt/go.py
+++ b/genopt/go.py
@@ -1,28 +1,35 @@
+import logging
+from typing import Callable, Tuple
+import sys
+
 import numpy as np
 
 from genopt.crossover import one_way_crossover
-from genopt.evaluation import decode, evaluate_pop
-from genopt.mutation import random_mutation
-from genopt.population import random_init, update_population
+from genopt.evaluation import decode_discrete, evaluate
+from genopt.mutation import mutation_discrete, mutation_real
+from genopt.population import init_discrete, init_real, update_population
 from genopt.selection import tournament_selection
 
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"
+)
+LOGGER = logging.getLogger(__name__)
 
 # pylint: disable=R0902
 class GeneticOptimizer:
     def __init__(
         self,
-        n_vars,
-        popsize,
-        n_gen,
-        objective_function,
-        t_sel_p=0.7,
-        t_sel_size=1,
-        mut_p=None,
-        mut_var=1,
-        encoding="real",
-        var_range=(0, 1),
-        var_size=1,
-        elitism=1,
+        n_vars: int,
+        popsize: int,
+        objective_function: Callable,
+        t_sel_p: float = 0.7,
+        t_sel_size: int = 1,
+        mut_p: float = None,
+        mut_var: float = 1.0,
+        encoding: str = "real",
+        var_range: Tuple[float] = (0, 1),
+        var_size: int = 1,
+        elitism: int = 1,
     ):
         """Maximize the function ``objective_function`` using genetic optimization.
         Supports real and discrete encoded variables.
@@ -31,21 +38,20 @@ class GeneticOptimizer:
             n_vars (int): Number of variables
             popsize (int): Population size, i.e number of chromosomes.
                 Needs to be an even number.
-            n_gen (int): Number of generations for the optimization to run
             objective_function (callable): Function to optimize for. Takes
                 the decoded variables as input and returns a fitness. The
                 genetic optimizer will then try to maximize the fitness.
             t_sel_p (float, optional): Probability of the fittest individual
                 to win a tournament. Defaults to 0.7.
             t_sel_size (int, optional): Tournament size. Defaults to 1.
-            mut_p ([type], optional): Mutation probability, set to
+            mut_p (float, optional): Mutation probability, set to
                 1/(n_vars*var_size) if None. Defaults to None.
             mut_var (int, optional): Mutation variance for real encoded
                 variables. Does nothing if encoding='discrete'. Defaults to 1.
             encoding (str, optional): Type of variable encoding. Can be 'real'
                 or 'discrete'. Defaults to 'real'.
             var_range (tuple, optional): Range of discrete variables. Does
-                nothing if encoding='real'. Defaults to (0,1).
+                nothing if encoding='real'. Defaults to (0, 1).
             var_size (int, optional): Number of genes for each discrete
                 variables. Multiplicatively increases chromosome length.
                 Does nothing if encoding='real'. Defaults to 1.
@@ -55,7 +61,6 @@ class GeneticOptimizer:
         """
 
         # Assertions
-        assert popsize % 2 == 0, "Popsize should be an even number"
         assert encoding in [
             "discrete",
             "real",
@@ -69,8 +74,6 @@ class GeneticOptimizer:
             ), "Variable size must be an integer larger than 0"
 
         self.n_vars = n_vars
-        self.popsize = popsize
-        self.n_gen = n_gen
         self.objective_function = objective_function
         self.t_sel_p = t_sel_p
         self.t_sel_size = t_sel_size
@@ -86,19 +89,30 @@ class GeneticOptimizer:
             self.mut_p = mut_p
         self.mut_var = mut_var
         self.elitism = elitism
-        self.fitness = np.zeros(self.popsize)
+        self.fitness = np.zeros(popsize)
         self.top_individual = None
 
         # Initialize population
-        self.population = random_init(
-            self.n_vars, self.popsize, self.encoding, self.var_size
-        )
+        if encoding == "real":
+            self.population = init_real(popsize, self.n_vars)
+        else:
+            self.population = init_discrete(popsize, self.n_vars, self.var_size)
 
-    def optimize(self):
-        for i in range(self.n_gen):
+    def optimize(self, n_gen: int) -> np.ndarray:
+        """Run the genetic optimizer for n_gen generations and return the
+        top individual of the population
+
+        Args:
+            n_gen (int): Number of generations to optimize for
+
+        Returns:
+            np.ndarray: Chromosome of the top individual
+        """
+        for i in range(n_gen):
 
             # Evaluate population
-            self.fitness, i_max = self.evaluate_pop(self.population)
+            self.fitness = self.evaluate(self.population)
+            i_max = np.argmax(self.fitness)
             self.top_individual = self.population[i_max, :]
 
             # Selection
@@ -114,7 +128,7 @@ class GeneticOptimizer:
             self.population = update_population(
                 tmp_population, self.top_individual, self.elitism
             )
-            print("Generation:", i, " - Max fitness:", self.fitness.max())
+            LOGGER.info(f"Generation: {i} - Max fitness: {self.fitness.max()}")
         return self.top_individual
 
     def decode(self, population: np.ndarray) -> np.ndarray:
@@ -127,16 +141,19 @@ class GeneticOptimizer:
         Returns:
             np.ndarray: Decoded population as 2d array of shape (popsize, n_vars)
         """
-        return decode(
+        if self.encoding == "real":
+            return population
+
+        return decode_discrete(
             population,
             self.n_vars,
-            self.encoding,
             self.var_range,
             self.var_size,
         )
 
-    def evaluate_pop(self, population: np.ndarray) -> np.ndarray:
-        """Return population fitness using the objective function. Decodes if necessary.
+    def evaluate(self, population: np.ndarray) -> np.ndarray:
+        """Evaluate population fitness using the objective function.
+        Decodes if necessary.
 
         Args:
             population (np.ndarray): Population as a 2d array of shape
@@ -145,13 +162,10 @@ class GeneticOptimizer:
         Returns:
             np.ndarray: Population fitness
         """
-        return evaluate_pop(
-            population,
-            self.n_vars,
+        variables = self.decode(population)
+        return evaluate(
+            variables,
             self.objective_function,
-            self.var_range,
-            self.var_size,
-            self.encoding,
         )
 
     def select(self, population: np.ndarray) -> np.ndarray:
@@ -165,10 +179,11 @@ class GeneticOptimizer:
             np.ndarray: Population after selection
         """
         return tournament_selection(
-            population, self.fitness, self.popsize, self.t_sel_p, self.t_sel_size
+            population, self.fitness, self.t_sel_p, self.t_sel_size
         )
 
-    def crossover(self, population: np.ndarray) -> np.ndarray:
+    @staticmethod
+    def crossover(population: np.ndarray) -> np.ndarray:
         """Mix the chromosomes of the population
 
         Args:
@@ -178,7 +193,7 @@ class GeneticOptimizer:
         Returns:
             np.ndarray: Population after crossover
         """
-        return one_way_crossover(population, self.n_vars, self.popsize)
+        return one_way_crossover(population)
 
     def mutate(self, population: np.ndarray) -> np.ndarray:
         """Mutate the population to introduce new chromosomes to the pool
@@ -190,12 +205,12 @@ class GeneticOptimizer:
         Returns:
             np.ndarray: Population after mutation
         """
-        return random_mutation(
+        if self.encoding == "real":
+            return mutation_real(population, self.n_vars, self.mut_p, self.mut_var)
+
+        return mutation_discrete(
             population,
             self.n_vars,
-            self.popsize,
             self.mut_p,
-            self.mut_var,
-            self.encoding,
             self.var_size,
         )
diff --git a/genopt/mutation.py b/genopt/mutation.py
index f615d51..41937da 100644
--- a/genopt/mutation.py
+++ b/genopt/mutation.py
@@ -1,21 +1,34 @@
 import numpy as np
 
 
-def random_mutation(
-    population,
-    n_vars,
-    popsize,
-    mut_p,
-    mut_var,
-    encoding,
-    var_size,
-):
-    if encoding == "real":
-        selected = np.random.rand(popsize, n_vars) < mut_p
-        population[selected] = population[selected] + np.random.normal(
-            0, mut_var, size=population[selected].shape
-        )
-    else:
-        selected = np.random.rand(popsize, n_vars * var_size) < mut_p
-        population[selected] = (population[selected] + 1) % 2
+def select_mutatees(
+    population: np.ndarray, chromosome_size: int, mut_p: float
+) -> np.ndarray:
+    popsize = population.shape[0]
+    return np.random.rand(popsize, chromosome_size) < mut_p
+
+
+def mutation_discrete(
+    population: np.ndarray,
+    n_vars: int,
+    mut_p: float,
+    var_size: int,
+) -> np.ndarray:
+    population = np.atleast_2d(population.copy())
+    selected = select_mutatees(population, n_vars * var_size, mut_p)
+    population[selected] = (population[selected] + 1) % 2
+    return population
+
+
+def mutation_real(
+    population: np.ndarray,
+    n_vars: int,
+    mut_p: float,
+    mut_var: float,
+) -> np.ndarray:
+    population = np.atleast_2d(population.copy())
+    selected = select_mutatees(population, n_vars, mut_p)
+    population[selected] = population[selected] + np.random.normal(
+        0, mut_var, size=population[selected].shape
+    )
     return population
diff --git a/genopt/population.py b/genopt/population.py
index 324eefb..3672179 100644
--- a/genopt/population.py
+++ b/genopt/population.py
@@ -1,14 +1,16 @@
 import numpy as np
 
 
-def random_init(n_vars, popsize, encoding, var_size):
-    if encoding == "real":
-        population = np.random.normal(0, 1, (popsize, n_vars))
-    else:
-        population = np.random.randint(2, size=(popsize, n_vars * var_size))
-    return population
+def init_discrete(popsize: int, n_vars: int, var_size: int):
+    return np.random.randint(2, size=(popsize, n_vars * var_size))
 
 
-def update_population(population, top_individual, elitism):
-    population[:elitism, :] = top_individual
-    return population
+def init_real(popsize: int, n_vars: int):
+    return np.random.normal(0, 1, (popsize, n_vars))
+
+
+def update_population(population: np.ndarray, top_individual: np.ndarray, elitism: int):
+    new_population = population.copy()
+    np.random.shuffle(new_population)
+    new_population[:elitism, :] = top_individual
+    return new_population
diff --git a/genopt/selection.py b/genopt/selection.py
index a7b9c6b..90ecfba 100644
--- a/genopt/selection.py
+++ b/genopt/selection.py
@@ -1,19 +1,33 @@
 import numpy as np
 
 
-def tournament_selection(population, fitness, popsize, t_sel_p, t_sel_size):
-    # Couldn't figure out a way to make it less messy without despicable for-loops
-    k = t_sel_size
+def tournament_selection(
+    population: np.ndarray, fitness: np.ndarray, t_sel_p: float, t_sel_size: int
+):
+    # Couldn't figure out a way to make it less messy without a bunch of loops
+    popsize = population.shape[0]
 
-    # Select k random individuals for each tournament
-    selected = np.random.randint(popsize, size=(popsize, k))
+    # Select t_sel_size random individuals for each tournament without replacement
+    selected = np.array(
+        [
+            np.random.choice(popsize, size=t_sel_size, replace=False)
+            for _ in range(popsize)
+        ]
+    )
     # Save their sorted indices
     i_sel = np.argsort(-np.take(fitness, selected), axis=1)
 
     # Select the best in t_sel_p of cases, nr 2 in t_sel_p*(1-t_sel_p)^1 of cases etc
-    probs = np.array([t_sel_p * (1 - t_sel_p) ** i for i in range(k - 1)])
-    probs = np.tile(np.append(np.cumsum(probs), 1), popsize).reshape((-1, k))
-    i_winner = np.argmax(np.random.rand(popsize).reshape((-1, 1)) < probs, axis=1)
+    prob_thresholds = np.cumsum(
+        [t_sel_p * (1 - t_sel_p) ** i for i in range(t_sel_size - 1)]
+    )
+    # Reshape probability thresholds so they can be broadcasted
+    prob_thresholds = np.tile(np.append(prob_thresholds, 1), popsize).reshape(
+        (-1, t_sel_size)
+    )
+    # Sample a random number in [0, 1] and see the threshold it lands in
+    prob_diffs = prob_thresholds - np.random.rand(popsize).reshape((-1, 1))
+    i_winner = np.argmax(np.minimum(prob_diffs, 0), axis=1)
 
     # Return the new population
     i_sel = i_sel[np.arange(popsize), i_winner]


Test patch:
diff --git a/tests/crossover_test.py b/tests/crossover_test.py
index 4640904..93db607 100644
--- a/tests/crossover_test.py
+++ b/tests/crossover_test.py
@@ -1,1 +1,28 @@
-# TODO
+import numpy as np
+from genopt.crossover import one_way_crossover
+
+
+def test_one_way_crossover_single_chromosome():
+    n_vars = 4
+    chromosome = np.random.rand(n_vars)
+    new_chromosome = one_way_crossover(chromosome)
+    assert (new_chromosome == chromosome).all()
+
+
+def test_one_way_crossover_even_popsize():
+    n_vars = 5
+    popsize = 2
+    population = np.random.rand(popsize, n_vars)
+    new_population = one_way_crossover(population)
+    assert population[0, 0] == new_population[1, 0]
+    assert population[1, 0] == new_population[0, 0]
+
+
+def test_one_way_crossover_uneven_popsize():
+    n_vars = 5
+    popsize = 3
+    population = np.random.rand(popsize, n_vars)
+    new_population = one_way_crossover(population)
+    assert population[0, 0] == new_population[1, 0]
+    assert population[1, 0] == new_population[0, 0]
+    assert (population[2, :] == new_population[2, :]).all()
diff --git a/tests/evaluation_test.py b/tests/evaluation_test.py
index 4640904..5682733 100644
--- a/tests/evaluation_test.py
+++ b/tests/evaluation_test.py
@@ -1,1 +1,52 @@
-# TODO
+import numpy as np
+
+from genopt.evaluation import evaluate, decode_discrete
+
+
+def objective(arr):
+    return arr.sum()
+
+
+def test_decode_discrete_single_chromosome():
+    n_vars = 10
+    var_size = 4
+    chromosome = np.random.randint(2, size=n_vars * var_size)
+    x = decode_discrete(chromosome, n_vars, (-1, 1), var_size)
+    assert np.logical_and(x >= -1, x <= 1).all()
+
+
+def test_decode_discrete_population():
+    n_vars = 10
+    var_size = 4
+    popsize = 10
+    chromosome = np.random.randint(2, size=(popsize, n_vars * var_size))
+    x = decode_discrete(chromosome, n_vars, (-1, 1), var_size)
+    assert np.logical_and(x >= -1, x <= 1).all()
+
+
+def test_decode_discrete_special_case():
+    n_vars = 10
+    var_size = 1
+    chromosome = np.random.randint(2, size=n_vars * var_size)
+    x = decode_discrete(chromosome, n_vars, (0, 1), var_size)
+    assert np.logical_and(x >= -1, x <= 1).all()
+    assert x is chromosome
+
+
+def test_evaluate_single_chromosome():
+    n_vars = 10
+    variables = np.random.rand(n_vars)
+    fitness = evaluate(variables, objective)
+    i_max = np.argmax(fitness)
+    assert fitness[0] == objective(variables)
+    assert i_max == 0
+
+
+def test_evaluate_population():
+    n_vars = 10
+    popsize = 25
+    variables = np.random.rand(popsize, n_vars)
+    variables[-1, :] = 1
+    fitness = evaluate(variables, objective)
+    i_max = np.argmax(fitness)
+    assert i_max == popsize - 1
diff --git a/tests/go_test.py b/tests/go_test.py
index ced68b0..cb4ba08 100644
--- a/tests/go_test.py
+++ b/tests/go_test.py
@@ -9,8 +9,8 @@ def objective(arr):
 def test_go_real():
     n_vars = 10
     random_result = objective(np.random.rand(n_vars))
-    go = GeneticOptimizer(n_vars, 100, 10, objective)
-    result = go.optimize()
+    go = GeneticOptimizer(n_vars, 100, objective)
+    result = go.optimize(10)
     assert objective(result) > random_result
 
 
@@ -19,8 +19,8 @@ def test_go_discrete():
     var_size = 4
     random_result = objective(np.random.rand(n_vars))
     go = GeneticOptimizer(
-        n_vars, 100, 10, objective, encoding="discrete", var_size=var_size
+        n_vars, 100, objective, encoding="discrete", var_size=var_size
     )
-    result = go.optimize()
+    result = go.optimize(10)
     result_decoded = go.decode(result)
     assert objective(result_decoded) > random_result
diff --git a/tests/mutation_test.py b/tests/mutation_test.py
index 4640904..2b076a4 100644
--- a/tests/mutation_test.py
+++ b/tests/mutation_test.py
@@ -1,1 +1,36 @@
-# TODO
+import numpy as np
+from genopt.mutation import mutation_discrete, mutation_real
+
+
+def test_mutation_real_single_chromosome():
+    n_vars = 10
+    chromosome = np.random.rand(n_vars)
+    mutated = mutation_real(chromosome.copy(), n_vars, 1, 1)
+    assert not (chromosome == mutated).any()
+
+
+def test_mutation_real():
+    n_vars = 10
+    popsize = 25
+    chromosome = np.random.rand(popsize, n_vars)
+    mutated = mutation_real(chromosome.copy(), n_vars, 0.5, 1)
+    assert not (chromosome == mutated).all()
+    assert (chromosome == mutated).any()
+
+
+def test_mutation_discrete_single_chromosome():
+    n_vars = 10
+    var_size = 4
+    chromosome = np.random.randint(2, size=n_vars * var_size)
+    mutated = mutation_discrete(chromosome.copy(), n_vars, 1, var_size)
+    assert not (chromosome == mutated).any()
+
+
+def test_mutation_discrete():
+    n_vars = 10
+    var_size = 4
+    popsize = 25
+    chromosome = np.random.randint(2, size=(popsize, n_vars * var_size))
+    mutated = mutation_discrete(chromosome.copy(), n_vars, 0.5, var_size)
+    assert not (chromosome == mutated).all()
+    assert (chromosome == mutated).any()
diff --git a/tests/population_test.py b/tests/population_test.py
index 4640904..678e637 100644
--- a/tests/population_test.py
+++ b/tests/population_test.py
@@ -1,1 +1,25 @@
-# TODO
+import numpy as np
+from genopt.population import init_discrete, init_real, update_population
+
+
+def test_init_discrete():
+    population = init_discrete(25, 10, 4)
+    assert population.shape == (25, 40)
+    assert np.logical_and(population >= 0, population <= 1).all()
+
+
+def test_init_discrete():
+    population = init_real(25, 10)
+    assert population.shape == (25, 10)
+    assert -0.5 <= population.mean() <= 0.5
+    assert 0.5 <= population.var() <= 1.5
+
+
+def test_update_population():
+    top_individual = np.ones(10)
+    population = init_real(25, 10)
+    new_population = update_population(population, top_individual, 2)
+    assert not (population == new_population).all()
+    assert (new_population[0, :] == top_individual).all()
+    assert (new_population[1, :] == top_individual).all()
+    assert (new_population[2, :] != top_individual).all()
diff --git a/tests/selection_test.py b/tests/selection_test.py
index 4640904..3d3a721 100644
--- a/tests/selection_test.py
+++ b/tests/selection_test.py
@@ -1,1 +1,48 @@
-# TODO
+from genopt.selection import tournament_selection
+import numpy as np
+from genopt.crossover import one_way_crossover
+
+
+def test_tournament_selection_size_2():
+    popsize = 2
+    t_sel_p = 0.7
+    t_sel_size = 2
+    fitness = np.arange(popsize)
+
+    wins = 0
+    iterations = 1000
+    for _ in range(iterations):
+        selected = tournament_selection(
+            fitness.reshape((-1, 1)), fitness, t_sel_p, t_sel_size
+        )
+        wins += (selected == 1).sum()
+
+    win_probability = wins / (popsize * iterations)
+    assert t_sel_p - 0.03 < win_probability < t_sel_p + 0.03
+
+
+def test_tournament_selection_size_5():
+    popsize = 5
+    t_sel_p = 0.7
+    t_sel_size = 5
+    fitness = np.arange(popsize)
+
+    wins = np.zeros(popsize)
+    iterations = 1000
+    for _ in range(iterations):
+        selected = tournament_selection(
+            fitness.reshape((-1, 1)), fitness, t_sel_p, t_sel_size
+        )
+
+        for i in range(popsize):
+            wins[i] += np.sum(selected == i)
+
+    wins /= iterations * t_sel_size
+
+    prob_thresholds = [t_sel_p * (1 - t_sel_p) ** i for i in range(t_sel_size - 1)]
+    prob_thresholds.append(1 - sum(prob_thresholds))
+
+    diff = 1e-2
+    for i, win in enumerate(wins):
+        prob = prob_thresholds[-(i + 1)]
+        assert prob - diff < win < prob + diff