From 436ba5c0e5352982e7298553e76dd98ba7464c3f Mon Sep 17 00:00:00 2001
From: Maximilian Schmidt <max.schmidt@fz-juelich.de>
Date: Fri, 20 Mar 2020 17:25:17 +0900
Subject: [PATCH] Add type hints in all files of the library

Start adding type hints

Black formatting and future imports to solve problem of forward reference

Fix more type hint issues and imports

Fix type annotations in population class and readd erroneously deleted variable in __init__ function

Fix some errors in node.py

Don't initialize class members

Make annotation more precise

Remove deprecated CGP in annotations

Fix initialization in genome

Updated hl_api: Removed Abstract
---
 gp/cartesian_graph.py |  57 ++++++++++----------
 gp/genome.py          |  95 +++++++++++++++++++++------------
 gp/hl_api.py          |  28 +++++-----
 gp/node.py            | 121 +++++++++++++++++++++++-------------------
 gp/population.py      |  34 +++++++-----
 gp/primitives.py      |  21 ++++----
 6 files changed, 202 insertions(+), 154 deletions(-)

diff --git a/gp/cartesian_graph.py b/gp/cartesian_graph.py
index 54981bd9..b4c9811d 100644
--- a/gp/cartesian_graph.py
+++ b/gp/cartesian_graph.py
@@ -13,6 +13,8 @@
 
 
 from .node import InputNode, OutputNode
+import gp
+from typing import Callable, List, Optional, Dict
 
 
 class CartesianGraph:
@@ -20,7 +22,7 @@ class CartesianGraph:
     Genome.
     """
 
-    def __init__(self, genome):
+    def __init__(self, genome: gp.Genome) -> None:
         """Init function.
 
         Parameters
@@ -28,35 +30,34 @@ def __init__(self, genome):
         genome: Genome
             Genome defining graph connectivity and node operations.
         """
-        self._n_outputs = None
-        self._n_inputs = None
-        self._n_columns = None
-        self._n_rows = None
-        self._nodes = None
-        self._gnome = None
+        self._n_outputs: int
+        self._n_inputs: int
+        self._n_columns: int
+        self._n_rows: int
+        self._nodes: List
 
         self.parse_genome(genome)
         self._genome = genome
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "CartesianGraph(" + str(self._nodes) + ")"
 
-    def print_active_nodes(self):
+    def print_active_nodes(self) -> str:
         """Print a representation of all active nodes in the graph.
         """
         return "CartesianGraph(" + str([node for node in self._nodes if node._active]) + ")"
 
-    def pretty_str(self):
+    def pretty_str(self) -> str:
         """Print a pretty representation of the Cartesian graph.
         """
         n_characters = 24
 
-        def pretty_node_str(node):
+        def pretty_node_str(node: gp.Node) -> str:
             s = node.pretty_str(n_characters)
             assert len(s) == n_characters
             return s
 
-        def empty_node_str():
+        def empty_node_str() -> str:
             return " " * n_characters
 
         s = "\n"
@@ -86,7 +87,7 @@ def empty_node_str():
 
         return s
 
-    def parse_genome(self, genome):
+    def parse_genome(self, genome: gp.Genome) -> None:
         if genome.dna is None:
             raise RuntimeError("dna not initialized")
 
@@ -114,22 +115,22 @@ def parse_genome(self, genome):
 
         self._determine_active_nodes()
 
-    def _hidden_column_idx(self, idx):
+    def _hidden_column_idx(self, idx: int) -> int:
         return (idx - self._n_inputs) // self._n_rows
 
     @property
-    def input_nodes(self):
+    def input_nodes(self) -> List[gp.Node]:
         return self._nodes[: self._n_inputs]
 
     @property
-    def hidden_nodes(self):
+    def hidden_nodes(self) -> List[gp.Node]:
         return self._nodes[self._n_inputs : -self._n_outputs]
 
     @property
-    def output_nodes(self):
+    def output_nodes(self) -> List[gp.Node]:
         return self._nodes[-self._n_outputs :]
 
-    def _determine_active_nodes(self):
+    def _determine_active_nodes(self) -> Dict[int, set]:
 
         # determine active nodes
         active_nodes_by_hidden_column_idx = collections.defaultdict(
@@ -155,7 +156,7 @@ def _determine_active_nodes(self):
 
         return active_nodes_by_hidden_column_idx
 
-    def determine_active_regions(self):
+    def determine_active_regions(self) -> List[int]:
         """Determine the active regions in the computational graph.
 
         Returns
@@ -171,7 +172,7 @@ def determine_active_regions(self):
 
         return active_regions
 
-    def __call__(self, x):
+    def __call__(self, x: List[float]) -> List[float]:
         # store values of x in input nodes
         for i, xi in enumerate(x):
             assert isinstance(self._nodes[i], InputNode)
@@ -185,16 +186,16 @@ def __call__(self, x):
 
         return [node._output for node in self.output_nodes]
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: int) -> gp.Node:
         return self._nodes[key]
 
-    def to_str(self):
+    def to_str(self) -> str:
 
         self._format_output_str_of_all_nodes()
         out_str = ", ".join(node.output_str for node in self.output_nodes)
         return f"[{out_str}]"
 
-    def _format_output_str_of_all_nodes(self):
+    def _format_output_str_of_all_nodes(self) -> None:
 
         for i, node in enumerate(self.input_nodes):
             node.format_output_str(self)
@@ -204,8 +205,8 @@ def _format_output_str_of_all_nodes(self):
             for node in active_nodes[hidden_column_idx]:
                 node.format_output_str(self)
 
-    def to_func(self):
-        """Compile the function(s) represented by the graph.
+    def to_func(self) -> Callable:
+        """Compile the function represented by the computational graph.
 
         Generates a definition of the function in Python code and
         executes the function definition to create a Callable.
@@ -270,7 +271,7 @@ def _f(x):
 
         return locals()["_f"]
 
-    def to_torch(self):
+    def to_torch(self) -> torch.nn.Module:
         """Compile the function(s) represented by the graph to a Torch class.
 
         Generates a definition of the Torch class in Python code and
@@ -347,7 +348,7 @@ def update_parameters_from_torch_class(self, torch_cls):
                 except AttributeError:
                     pass
 
-    def to_sympy(self, simplify=True):
+    def to_sympy(self, simplify: Optional[bool] = True) -> List[sympy.core.Expr]:
         """Compile the function(s) represented by the graph to a SymPy expression.
 
         Generates one SymPy expression for each output node.
@@ -366,7 +367,7 @@ def to_sympy(self, simplify=True):
         if sympy is None:
             raise ModuleNotFoundError("No module named 'sympy' (extra requirement)")
 
-        def _validate_sympy_expr(expr):
+        def _validate_sympy_expr(expr: sympy.core.Expr) -> sympy.core.Expr:
             """Helper function that raises an exception upon encountering a SymPy
             expression that can not be evaluated.
 
diff --git a/gp/genome.py b/gp/genome.py
index 988a0dbb..bf594876 100644
--- a/gp/genome.py
+++ b/gp/genome.py
@@ -1,3 +1,8 @@
+import gp
+import numpy as np
+
+from typing import Generator, List, Optional, Tuple
+
 from .primitives import Primitives
 
 
@@ -5,7 +10,15 @@ class Genome:
     """Genome class for  individuals.
     """
 
-    def __init__(self, n_inputs, n_outputs, n_columns, n_rows, levels_back, primitives):
+    def __init__(
+        self,
+        n_inputs: int,
+        n_outputs: int,
+        n_columns: int,
+        n_rows: int,
+        levels_back: int,
+        primitives: List[gp.Node],
+    ) -> None:
         """Init function.
 
         Parameters
@@ -21,7 +34,7 @@ def __init__(self, n_inputs, n_outputs, n_columns, n_rows, levels_back, primitiv
         levels_back : int
             Number of previous columns that an entry in the genome can be
             connected with.
-        primitives : List[gp.CPGNode]
+        primitives : List[gp.Node]
            List of primitives that the genome can refer to.
         """
         if n_inputs <= 0:
@@ -53,46 +66,46 @@ def __init__(self, n_inputs, n_outputs, n_columns, n_rows, levels_back, primitiv
             1 + self._primitives.max_arity
         )  # one function gene + multiple input genes
 
-        self._dna = None  # stores dna as list of alleles for all regions
+        self._dna: List[int] = []  # stores dna as list of alleles for all regions
 
         # constants used as identifiers for input and output nodes
-        self._id_input_node = -1
-        self._id_output_node = -2
-        self._non_coding_allele = None
+        self._id_input_node: int = -1
+        self._id_output_node: int = -2
+        self._non_coding_allele: int
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: int) -> int:
         if self._dna is None:
             raise RuntimeError("dna not initialized")
         return self._dna[key]
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: int, value: int) -> None:
         dna = list(self._dna)
         dna[key] = value
         self._validate_dna(dna)
         self._dna = dna
 
     @property
-    def dna(self):
+    def dna(self) -> List[int]:
         return self._dna
 
     @dna.setter
-    def dna(self, value):
+    def dna(self, value: List[int]) -> None:
         self._validate_dna(value)
         self._dna = value
 
     @property
-    def _n_hidden(self):
+    def _n_hidden(self) -> int:
         return self._n_columns * self._n_rows
 
     @property
-    def _n_regions(self):
+    def _n_regions(self) -> int:
         return self._n_inputs + self._n_hidden + self._n_outputs
 
     @property
-    def _n_genes(self):
+    def _n_genes(self) -> int:
         return self._n_regions * self._length_per_region
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         s = self.__class__.__name__ + "("
         for region_idx, input_region in self.iter_input_regions():
             s += str(region_idx) + ": " + str(input_region) + " | "
@@ -104,7 +117,7 @@ def __repr__(self):
         s += ")"
         return s
 
-    def _create_input_region(self):
+    def _create_input_region(self) -> list:
 
         # fill region with identifier for input node and zeros,
         # since input nodes do not have any inputs
@@ -113,7 +126,9 @@ def _create_input_region(self):
         region += [self._non_coding_allele] * self._primitives.max_arity
         return region
 
-    def _create_random_hidden_region(self, rng, permissable_inputs):
+    def _create_random_hidden_region(
+        self, rng: np.random.RandomState, permissable_inputs: List[int]
+    ) -> list:
 
         # construct dna region consisting of function allele and
         # input alleles
@@ -131,7 +146,9 @@ def _create_random_hidden_region(self, rng, permissable_inputs):
 
         return region
 
-    def _create_random_output_region(self, rng, permissable_inputs):
+    def _create_random_output_region(
+        self, rng: np.random.RandomState, permissable_inputs: List[int]
+    ) -> list:
 
         # fill region with identifier for output node and single
         # gene determining input
@@ -142,7 +159,7 @@ def _create_random_output_region(self, rng, permissable_inputs):
 
         return region
 
-    def randomize(self, rng):
+    def randomize(self, rng: np.random.RandomState) -> None:
         """Randomize the genome.
 
         Parameters
@@ -177,7 +194,7 @@ def randomize(self, rng):
         self._validate_dna(dna)
         self._dna = dna
 
-    def _permissable_inputs(self, region_idx):
+    def _permissable_inputs(self, region_idx: int) -> List[int]:
 
         assert not self._is_input_region(region_idx)
 
@@ -200,10 +217,10 @@ def _permissable_inputs(self, region_idx):
 
         return permissable_inputs
 
-    def _permissable_inputs_for_output_region(self):
+    def _permissable_inputs_for_output_region(self) -> list:
         return self._permissable_inputs(self._n_inputs + self._n_rows * self._n_columns)
 
-    def _validate_dna(self, dna):
+    def _validate_dna(self, dna: List[int]) -> None:
 
         if len(dna) != self._n_genes:
             raise ValueError("dna length mismatch")
@@ -257,7 +274,7 @@ def _validate_dna(self, dna):
                     "be identical to non-coding allele"
                 )
 
-    def _hidden_column_idx(self, region_idx):
+    def _hidden_column_idx(self, region_idx: int) -> int:
         assert self._n_inputs <= region_idx
         assert region_idx < self._n_inputs + self._n_rows * self._n_columns
         hidden_column_idx = (region_idx - self._n_inputs) // self._n_rows
@@ -265,7 +282,9 @@ def _hidden_column_idx(self, region_idx):
         assert hidden_column_idx < self._n_columns
         return hidden_column_idx
 
-    def iter_input_regions(self, dna=None):
+    def iter_input_regions(
+        self, dna: Optional[List[int]] = None
+    ) -> Generator[Tuple[int, list], None, None]:
         if dna is None:
             dna = self.dna
         for i in range(self._n_inputs):
@@ -275,7 +294,9 @@ def iter_input_regions(self, dna=None):
             ]
             yield region_idx, region
 
-    def iter_hidden_regions(self, dna=None):
+    def iter_hidden_regions(
+        self, dna: Optional[List[int]] = None
+    ) -> Generator[Tuple[int, List[int]], None, None]:
         if dna is None:
             dna = self.dna
         for i in range(self._n_hidden):
@@ -285,7 +306,9 @@ def iter_hidden_regions(self, dna=None):
             ]
             yield region_idx, region
 
-    def iter_output_regions(self, dna=None):
+    def iter_output_regions(
+        self, dna: Optional[List[int]] = None
+    ) -> Generator[Tuple[int, List[int]], None, None]:
         if dna is None:
             dna = self.dna
         for i in range(self._n_outputs):
@@ -295,22 +318,22 @@ def iter_output_regions(self, dna=None):
             ]
             yield region_idx, region
 
-    def _is_input_region(self, region_idx):
+    def _is_input_region(self, region_idx: int) -> bool:
         return region_idx < self._n_inputs
 
-    def _is_hidden_region(self, region_idx):
+    def _is_hidden_region(self, region_idx: int) -> bool:
         return (self._n_inputs <= region_idx) & (region_idx < self._n_inputs + self._n_hidden)
 
-    def _is_output_region(self, region_idx):
+    def _is_output_region(self, region_idx: int) -> bool:
         return self._n_inputs + self._n_hidden <= region_idx
 
-    def _is_function_gene(self, idx):
+    def _is_function_gene(self, idx: int) -> bool:
         return (idx % self._length_per_region) == 0
 
-    def _is_input_gene(self, idx):
+    def _is_input_gene(self, idx: int) -> bool:
         return not self._is_function_gene(idx)
 
-    def mutate(self, n_mutations, active_regions, rng):
+    def mutate(self, n_mutations: int, active_regions: List[int], rng: np.random.RandomState):
         """Mutate the genome.
 
         Parameters
@@ -321,7 +344,7 @@ def mutate(self, n_mutations, active_regions, rng):
             Regions in the genome that are currently used in the
             computational graph. Used to check whether mutations are
             silent or require reevaluation of fitness.
-        rng : numpy.RandomState
+        rng : numpy.random.RandomState
             Random number generator instance to use for crossover.
 
         Returns
@@ -373,7 +396,9 @@ def _mutate_output_region(self, gene_idx, region_idx, rng):
 
         return False
 
-    def _mutate_hidden_region(self, gene_idx, region_idx, rng):
+    def _mutate_hidden_region(
+        self, gene_idx: int, region_idx: int, rng: np.random.RandomState
+    ) -> bool:
         assert self._is_hidden_region(region_idx)
 
         if self._is_function_gene(gene_idx):
@@ -406,10 +431,10 @@ def _mutate_hidden_region(self, gene_idx, region_idx, rng):
         return False
 
     @property
-    def primitives(self):
+    def primitives(self) -> gp.Primitives:
         return self._primitives
 
-    def clone(self):
+    def clone(self) -> gp.Genome:
         """Clone the genome.
 
         Returns
diff --git a/gp/hl_api.py b/gp/hl_api.py
index aa3ea934..6fc18383 100644
--- a/gp/hl_api.py
+++ b/gp/hl_api.py
@@ -1,18 +1,21 @@
 import numpy as np
+from typing import Optional, Callable
+from .population import Population
+from .individual import Individual
+from .ea import MuPlusLambda
 
 
 def evolve(
-    pop,
-    objective,
-    ea,
-    max_generations,
-    min_fitness,
-    print_progress=False,
-    *,
-    callback=None,
-    label=None,
-    n_processes=1,
-):
+    pop: Population,
+    objective: Callable[[Individual], Individual],
+    ea: MuPlusLambda,
+    max_generations: int,
+    min_fitness: float,
+    print_progress: Optional[bool] = False,
+    callback: Optional[Callable[[Population], None]] = None,
+    label: Optional[str] = None,
+    n_processes: int = 1,
+) -> None:
     """
     Evolves a population and returns the history of fitness of parents.
 
@@ -45,8 +48,7 @@ def evolve(
 
     Returns
     -------
-    dict
-        History of the evolution.
+    None
     """
 
     ea.initialize_fitness_parents(pop, objective, label=label)
diff --git a/gp/node.py b/gp/node.py
index e81f32c6..6ab113d2 100644
--- a/gp/node.py
+++ b/gp/node.py
@@ -1,12 +1,17 @@
+from __future__ import annotations
+
+from typing import List, Union
+from cartesian_graph import CartesianGraph
+
 primitives_dict = {}  # maps string of class names to classes
 
 
-def register(cls):
+def register(cls: Node) -> None:
     """Register a primitive in the global dictionary of primitives
 
     Parameters
     ----------
-    cls : gp.CPGNode
+    cls : gp.Node
        Primitive to be registered.
 
     Returns
@@ -22,14 +27,17 @@ class Node:
     """Base class for primitive functions used in Cartesian computational graphs.
     """
 
-    _arity = None
-    _active = False
-    _inputs = None
-    _output = None
-    _idx = None
-    _is_parameter = False
-
-    def __init__(self, idx, inputs):
+    __name__ = "Node"
+    _arity: int
+    _active: bool
+    _inputs: List[int]
+    _output: float
+    _idx: int
+    _is_parameter: bool
+    _output_str: str
+    _parameter_str: str
+
+    def __init__(self, idx: int, inputs: List[int]) -> None:
         """Init function.
 
         Parameters
@@ -44,33 +52,36 @@ def __init__(self, idx, inputs):
 
         assert idx not in inputs
 
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
-        register(cls)
+    # def __init_subclass__(cls: Node, **kwargs: Any) -> None:
+    #     super().__init_subclass__(**kwargs)
+    #     register(cls)
+
+    def __call__(self, x: float, graph: CartesianGraph) -> None:
+        raise NotImplementedError
 
     @property
-    def arity(self):
+    def arity(self) -> Union[None, int]:
         return self._arity
 
     @property
-    def max_arity(self):
+    def max_arity(self) -> int:
         return len(self._inputs)
 
     @property
-    def inputs(self):
+    def inputs(self) -> List[int]:
         return self._inputs[: self._arity]
 
     @property
-    def idx(self):
+    def idx(self) -> Union[None, int]:
         return self._idx
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__}(idx: {self.idx}, active: {self._active}, "
             f"arity: {self._arity}, inputs {self._inputs}, output {self._output})"
         )
 
-    def pretty_str(self, n):
+    def pretty_str(self, n: int) -> str:
         used_characters = 0
         used_characters += 3  # for two digit idx and following whitespace
         used_characters += 3  # for "active" marker
@@ -111,15 +122,15 @@ def pretty_str(self, n):
         return s.ljust(n)
 
     @property
-    def output(self):
+    def output(self) -> float:
         return self._output
 
-    def activate(self):
+    def activate(self) -> None:
         """Set node to active.
         """
         self._active = True
 
-    def format_output_str(self, graph):
+    def format_output_str(self, graph: CartesianGraph) -> None:
         """Format output string of the node.
         """
         raise NotImplementedError()
@@ -140,19 +151,19 @@ def format_output_str_torch(self, graph):
         """
         self.format_output_str(graph)
 
-    def format_parameter_str(self):
+    def format_parameter_str(self) -> None:
         raise NotImplementedError
 
     @property
-    def output_str(self):
+    def output_str(self) -> str:
         return self._output_str
 
     @property
-    def is_parameter(self):
+    def is_parameter(self) -> bool:
         return self._is_parameter
 
     @property
-    def parameter_str(self):
+    def parameter_str(self) -> str:
         return self._parameter_str
 
 
@@ -162,13 +173,13 @@ class Add(Node):
 
     _arity = 2
 
-    def __init__(self, idx, inputs):
+    def __init__(self, idx: int, inputs: List[int]) -> None:
         super().__init__(idx, inputs)
 
-    def __call__(self, x, graph):
+    def __call__(self, x: float, graph: CartesianGraph) -> None:
         self._output = graph[self._inputs[0]].output + graph[self._inputs[1]].output
 
-    def format_output_str(self, graph):
+    def format_output_str(self, graph: CartesianGraph) -> None:
         self._output_str = (
             f"({graph[self._inputs[0]].output_str} + {graph[self._inputs[1]].output_str})"
         )
@@ -180,13 +191,13 @@ class Sub(Node):
 
     _arity = 2
 
-    def __init__(self, idx, inputs):
+    def __init__(self, idx: int, inputs: List[int]) -> None:
         super().__init__(idx, inputs)
 
-    def __call__(self, x, graph):
+    def __call__(self, x: float, graph: CartesianGraph) -> None:
         self._output = graph[self._inputs[0]].output - graph[self._inputs[1]].output
 
-    def format_output_str(self, graph):
+    def format_output_str(self, graph: CartesianGraph) -> None:
         self._output_str = (
             f"({graph[self._inputs[0]].output_str} - {graph[self._inputs[1]].output_str})"
         )
@@ -198,13 +209,13 @@ class Mul(Node):
 
     _arity = 2
 
-    def __init__(self, idx, inputs):
+    def __init__(self, idx: int, inputs: List[int]) -> None:
         super().__init__(idx, inputs)
 
-    def __call__(self, x, graph):
+    def __call__(self, x: float, graph: CartesianGraph) -> None:
         self._output = graph[self._inputs[0]].output * graph[self._inputs[1]].output
 
-    def format_output_str(self, graph):
+    def format_output_str(self, graph: CartesianGraph) -> None:
         self._output_str = (
             f"({graph[self._inputs[0]].output_str} * {graph[self._inputs[1]].output_str})"
         )
@@ -216,14 +227,14 @@ class Div(Node):
 
     _arity = 2
 
-    def __init__(self, idx, inputs):
+    def __init__(self, idx: int, inputs: List[int]) -> None:
         super().__init__(idx, inputs)
 
-    def __call__(self, x, graph):
+    def __call__(self, x: float, graph: CartesianGraph) -> None:
 
         self._output = graph[self._inputs[0]].output / graph[self._inputs[1]].output
 
-    def format_output_str(self, graph):
+    def format_output_str(self, graph: CartesianGraph) -> None:
         self._output_str = (
             f"({graph[self._inputs[0]].output_str} / {graph[self._inputs[1]].output_str})"
         )
@@ -236,24 +247,24 @@ class ConstantFloat(Node):
     _arity = 0
     _is_parameter = True
 
-    def __init__(self, idx, inputs):
+    def __init__(self, idx: int, inputs: List[int]) -> None:
         super().__init__(idx, inputs)
 
         self._output = 1.0
 
-    def __call__(self, x, graph):
+    def __call__(self, x: float, graph: CartesianGraph) -> None:
         pass
 
-    def format_output_str(self, graph):
+    def format_output_str(self, graph: CartesianGraph) -> None:
         self._output_str = f"{self._output}"
 
-    def format_output_str_numpy(self, graph):
+    def format_output_str_numpy(self, graph: CartesianGraph) -> None:
         self._output_str = f"np.ones(x.shape[0]) * {self._output}"
 
-    def format_output_str_torch(self, graph):
+    def format_output_str_torch(self, graph: CartesianGraph) -> None:
         self._output_str = f"self._p{self._idx}.expand(x.shape[0])"
 
-    def format_parameter_str(self):
+    def format_parameter_str(self) -> None:
         self._parameter_str = (
             f"self._p{self._idx} = torch.nn.Parameter(torch.Tensor([{self._output}]))\n"
         )
@@ -265,19 +276,19 @@ class InputNode(Node):
 
     _arity = 0
 
-    def __init__(self, idx, inputs):
+    def __init__(self, idx: int, inputs: List[int]) -> None:
         super().__init__(idx, inputs)
 
-    def __call__(self, x, graph):
+    def __call__(self, x: float, graph: CartesianGraph) -> None:
         assert False
 
-    def format_output_str(self, graph):
+    def format_output_str(self, graph: CartesianGraph) -> None:
         self._output_str = f"x[{self._idx}]"
 
-    def format_output_str_numpy(self, graph):
+    def format_output_str_numpy(self, graph: CartesianGraph) -> None:
         self._output_str = f"x[:, {self._idx}]"
 
-    def format_output_str_torch(self, graph):
+    def format_output_str_torch(self, graph: CartesianGraph) -> None:
         self._output_str = f"x[:, {self._idx}]"
 
 
@@ -287,13 +298,13 @@ class OutputNode(Node):
 
     _arity = 1
 
-    def __init__(self, idx, inputs):
+    def __init__(self, idx: int, inputs: List[int]) -> None:
         super().__init__(idx, inputs)
 
-    def __call__(self, x, graph):
+    def __call__(self, x: float, graph: CartesianGraph) -> None:
         self._output = graph[self._inputs[0]].output
 
-    def format_output_str(self, graph):
+    def format_output_str(self, graph: CartesianGraph) -> None:
         self._output_str = f"{graph[self._inputs[0]].output_str}"
 
 
@@ -303,13 +314,13 @@ class Pow(Node):
 
     _arity = 2
 
-    def __init__(self, idx, inputs):
+    def __init__(self, idx: int, inputs: List[int]) -> None:
         super().__init__(idx, inputs)
 
-    def __call__(self, x, graph):
+    def __call__(self, x: float, graph: CartesianGraph) -> None:
         self._output = graph[self._inputs[0]].output ** graph[self._inputs[1]].output
 
-    def format_output_str(self, graph):
+    def format_output_str(self, graph: CartesianGraph) -> None:
         self._output_str = (
             f"({graph[self._inputs[0]].output_str} ** {graph[self._inputs[1]].output_str})"
         )
diff --git a/gp/population.py b/gp/population.py
index 2a2d8722..4586b80e 100755
--- a/gp/population.py
+++ b/gp/population.py
@@ -1,6 +1,8 @@
+import gp
 import numpy as np
 
 from .individual import Individual, IndividualMultiGenome
+from typing import List
 
 
 class Population:
@@ -8,7 +10,9 @@ class Population:
     A population of individuals.
     """
 
-    def __init__(self, n_parents, mutation_rate, seed, genome_params):
+    def __init__(
+        self, n_parents: int, mutation_rate: float, seed: int, genome_params: dict
+    ) -> None:
         """Init function.
 
         Parameters
@@ -34,7 +38,7 @@ def __init__(self, n_parents, mutation_rate, seed, genome_params):
 
         self._genome_params = genome_params
 
-        self._parents = None  # list of parent individuals
+        self._parents: List[Individual] = []  # list of parent individuals
 
         # keeps track of the number of generations, increases with
         # every new offspring generation
@@ -44,33 +48,33 @@ def __init__(self, n_parents, mutation_rate, seed, genome_params):
         self._generate_random_parent_population()
 
     @property
-    def champion(self):
+    def champion(self) -> gp.Individual:
         """Return parent with the highest fitness.
         """
         return max(self._parents, key=lambda ind: ind.fitness)
 
     @property
-    def parents(self):
+    def parents(self) -> List[gp.Individual]:
         return self._parents
 
     @parents.setter
-    def parents(self, new_parents):
+    def parents(self, new_parents: List[gp.Individual]) -> None:
         self.generation += 1
         self._parents = new_parents
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> gp.Individual:
         return self._parents[idx]
 
-    def _generate_random_parent_population(self):
+    def _generate_random_parent_population(self) -> None:
         self._parents = self._generate_random_individuals(self.n_parents)
         self._label_new_individuals(self._parents)
 
-    def _label_new_individuals(self, individuals):
+    def _label_new_individuals(self, individuals: List[gp.Individual]) -> None:
         for ind in individuals:
             ind.idx = self._max_idx + 1
             self._max_idx += 1
 
-    def _generate_random_individuals(self, n):
+    def _generate_random_individuals(self, n: int) -> List[Individual]:
         individuals = []
         for i in range(n):
             if isinstance(self._genome_params, dict):
@@ -86,12 +90,14 @@ def _generate_random_individuals(self, n):
 
         return individuals
 
-    def crossover(self, breeding_pool, n_offsprings):
+    def crossover(
+        self, breeding_pool: List[gp.Individual], n_offsprings: int
+    ) -> List[gp.Individual]:
         """Create an offspring population via crossover.
 
         Parameters
         ----------
-        breeding_pool : List[gp.Population]
+        breeding_pool : List[gp.Individual]
             List of individuals from which the offspring are created.
         n_offsprings : int
             Number of offspring to be created.
@@ -121,7 +127,7 @@ def crossover(self, breeding_pool, n_offsprings):
         assert len(breeding_pool) >= n_offsprings
         return sorted(breeding_pool, key=lambda x: -x.fitness)[:n_offsprings]
 
-    def mutate(self, offsprings):
+    def mutate(self, offsprings: List[gp.Individual]) -> List[gp.Individual]:
         """Mutate a list of offspring invididuals.
 
         Parameters
@@ -139,7 +145,7 @@ def mutate(self, offsprings):
             off.mutate(self._mutation_rate, self.rng)
         return offsprings
 
-    def fitness_parents(self):
+    def fitness_parents(self) -> List[float]:
         """Return fitness for all parents of the population.
 
         Returns
@@ -149,7 +155,7 @@ def fitness_parents(self):
         """
         return [ind.fitness for ind in self._parents]
 
-    def dna_parents(self):
+    def dna_parents(self) -> List[List[int]]:
         """Return a list of the DNA of all parents.
 
         Returns
diff --git a/gp/primitives.py b/gp/primitives.py
index d925a722..8c3482dd 100644
--- a/gp/primitives.py
+++ b/gp/primitives.py
@@ -1,5 +1,8 @@
 from .node import Node
 
+from typing import List, Tuple
+import numpy as np  # type: ignore
+
 
 class Primitives:
     """Class collecting primitives of the Cartesian Genetic Programming framework.
@@ -7,14 +10,14 @@ class Primitives:
 
     _n_primitives = 0
     _max_arity = 0
-    _primitives = None
+    _primitives: dict = {}
 
-    def __init__(self, primitives):
+    def __init__(self, primitives: List[Node]) -> None:
         """Init function.
 
         Parameters
         ----------
-        primitives : List[gp.CPGNode]
+        primitives : List[gp.Node]
             List of primitives.
         """
         for i in range(len(primitives)):
@@ -37,7 +40,7 @@ def __init__(self, primitives):
 
         self._determine_max_arity()
 
-    def _determine_max_arity(self):
+    def _determine_max_arity(self) -> None:
 
         arity = 1  # minimal possible arity (output nodes need one input)
 
@@ -47,7 +50,7 @@ def _determine_max_arity(self):
 
         self._max_arity = arity
 
-    def sample(self, rng):
+    def sample(self, rng: np.random.RandomState) -> Node:
         """Sample a random primitive.
 
         Parameters
@@ -57,19 +60,19 @@ def sample(self, rng):
 
         Returns
         -------
-        List[Tuple(str, gp.CPGNode)]
+        List[Tuple(str, gp.Node)]
         """
         return rng.choice(self.alleles)
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: int) -> Node:
         return self._primitives[key]
 
     @property
-    def max_arity(self):
+    def max_arity(self) -> int:
         return self._max_arity
 
     @property
-    def alleles(self):
+    def alleles(self) -> Tuple:
         return tuple(self._primitives.keys())
 
     def __len__(self):