simplify function set creation

trevorstephens · Aug 27, 2016 · a3d0f9f · a3d0f9f
1 parent 65e0305
commit a3d0f9f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 15 deletions.
diff --git a/gplearn/functions.py b/gplearn/functions.py
@@ -146,17 +146,17 @@ def _protected_inverse(x1):
 cos1 = make_function(function=np.cos, name='cos', arity=1)
 tan1 = make_function(function=np.tan, name='tan', arity=1)
 
-_function_map = {'add': [add2],
-                 'sub': [sub2],
-                 'mul': [mul2],
-                 'div': [div2],
-                 'sqrt': [sqrt1],
-                 'log': [log1],
-                 'abs': [abs1],
-                 'neg': [neg1],
-                 'inv': [inv1],
-                 'max': [max2],
-                 'min': [min2],
-                 'sin': [sin1],
-                 'cos': [cos1],
-                 'tan': [tan1]}
+_function_map = {'add': add2,
+                 'sub': sub2,
+                 'mul': mul2,
+                 'div': div2,
+                 'sqrt': sqrt1,
+                 'log': log1,
+                 'abs': abs1,
+                 'neg': neg1,
+                 'inv': inv1,
+                 'max': max2,
+                 'min': min2,
+                 'sin': sin1,
+                 'cos': cos1,
+                 'tan': tan1}
diff --git a/gplearn/genetic.py b/gplearn/genetic.py
@@ -932,7 +932,7 @@ def fit(self, X, y, sample_weight=None):
                 if function not in _function_map:
                     raise ValueError('invalid function name %s found in '
                                      '`function_set`.' % function)
-                self._function_set.extend(_function_map[function])
+                self._function_set.append(_function_map[function])
             elif isinstance(function, _Function):
                 self._function_set.append(function)
             else: