Review fixes

aimclub · Nov 28, 2023 · acc7a89 · acc7a89
1 parent 620279d
commit acc7a89
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 23 deletions.
diff --git a/golem/core/optimisers/genetic/operators/mutation.py b/golem/core/optimisers/genetic/operators/mutation.py
@@ -118,30 +118,22 @@ def _mutation(self, individual: Individual) -> Tuple[Individual, bool]:
                                'Please check optimization parameters for correctness.')
         return individual, is_applied
 
-    def _sample_num_of_mutations(self) -> int:
+    def _sample_num_of_mutations(self, mutation_type: MutationIdType) -> int:
         # most of the time returns 1 or rarely several mutations
-        if self.parameters.variable_mutation_num:
+        is_custom_mutation = isinstance(mutation_type, Callable)
+        if self.parameters.variable_mutation_num and not is_custom_mutation:
             num_mut = max(int(round(np.random.lognormal(0, sigma=0.5))), 1)
         else:
             num_mut = 1
         return num_mut
 
     def _apply_mutations(self, new_graph: Graph, mutation_type: MutationIdType) -> Graph:
         """Apply mutation 1 or few times iteratively"""
-        for _ in range(self._sample_num_of_mutations()):
-            new_graph = self._adapt_and_apply_mutation(new_graph, mutation_type)
-
-            is_custom_mutation = isinstance(mutation_type, Callable)
-            if is_custom_mutation:  # custom mutation occurs once
-                break
-        return new_graph
-
-    def _adapt_and_apply_mutation(self, new_graph: Graph, mutation_type: MutationIdType) -> Graph:
-        # get the mutation function and adapt it
-        mutation_func = self._get_mutation_func(mutation_type)
-        new_graph = mutation_func(new_graph, requirements=self.requirements,
-                                  graph_gen_params=self.graph_generation_params,
-                                  parameters=self.parameters)
+        for _ in range(self._sample_num_of_mutations(mutation_type)):
+            mutation_func = self._get_mutation_func(mutation_type)
+            new_graph = mutation_func(new_graph, requirements=self.requirements,
+                                      graph_gen_params=self.graph_generation_params,
+                                      parameters=self.parameters)
         return new_graph
 
     def _will_mutation_be_applied(self, mutation_type: MutationIdType) -> bool:

diff --git a/test/integration/test_genetic_schemes.py b/test/integration/test_genetic_schemes.py
@@ -1,9 +1,8 @@
-from functools import partial
+from typing import Sequence
 
 import numpy as np
 import pytest
 
-from examples.synthetic_graph_evolution.experiment_setup import run_trial
 from examples.synthetic_graph_evolution.generators import generate_labeled_graph
 from examples.synthetic_graph_evolution.tree_search import tree_search_setup
 from golem.core.optimisers.genetic.gp_params import GPAlgorithmParameters
@@ -17,7 +16,7 @@ def set_up_params(genetic_scheme: GeneticSchemeTypesEnum):
         multi_objective=False,
         mutation_types=[
             MutationTypesEnum.single_add,
-            MutationTypesEnum.single_drop,
+            MutationTypesEnum.single_drop
         ],
         crossover_types=[CrossoverTypesEnum.none],
         genetic_scheme_type=genetic_scheme
@@ -27,13 +26,16 @@ def set_up_params(genetic_scheme: GeneticSchemeTypesEnum):
 
 @pytest.mark.parametrize('genetic_type', GeneticSchemeTypesEnum)
 def test_genetic_scheme_types(genetic_type):
-    target_graph = generate_labeled_graph('tree', 4, node_labels=['x'])
+    target_graph = generate_labeled_graph('tree', 30, node_labels=['x'])
     num_iterations = 30
 
     gp_params = set_up_params(genetic_type)
-    found_graph, history = run_trial(target_graph=target_graph,
-                                     optimizer_setup=partial(tree_search_setup, algorithm_parameters=gp_params),
-                                     num_iterations=num_iterations)
+    optimizer, objective = tree_search_setup(target_graph,
+                                             num_iterations=num_iterations,
+                                             algorithm_parameters=gp_params)
+    found_graphs = optimizer.optimise(objective)
+    found_graph = found_graphs[0] if isinstance(found_graphs, Sequence) else found_graphs
+    history = optimizer.history
     assert found_graph is not None
     # at least 20% more generation than early_stopping_iterations were evaluated
     # (+2 gen for initial assumption and final choice)