DataSaver doesn't work with the BalancingLearner

[jbweston]

(original issue on GitLab)

opened by Bas Nijholt (@basnijholt) at 2017-11-02T11:14:36.670Z

import adaptive
from adaptive.learner import DataSaver, Learner1D, BalancingLearner
from adaptive.runner import SequentialExecutor
from operator import itemgetter
adaptive.notebook_extension()

def f(x):
    return {'y': x**2, 'x': 'hi'}

_learners = [Learner1D(f, (-1, 1)) for _ in range(2)]
learner = DataSaver(BalancingLearner(_learners), arg_picker=itemgetter('y'))

points, _ = learner.choose_points(1)
for x in points:
    learner.add_point(x, learner.function(x))

error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-21-2018dbfa55fc> in <module>()
     11 learner = BalancingLearner(_learners)
     12 
---> 13 learner.choose_points(1)

/opt/conda/lib/python3.6/site-packages/adaptive/learner/balancing_learner.py in choose_points(self, n, add_data)
     65                 return self._choose_and_add_points(n)
     66         else:
---> 67             return self._choose_and_add_points(n)
     68 
     69     def add_point(self, x, y):

/opt/conda/lib/python3.6/site-packages/adaptive/learner/balancing_learner.py in _choose_and_add_points(self, n)
     56             x, _ = max(zip(pairs, loss_improvements), key=itemgetter(1))
     57             points.append(x)
---> 58             self.add_point(x, None)
     59         return points, None
     60 

/opt/conda/lib/python3.6/site-packages/adaptive/learner/balancing_learner.py in add_point(self, x, y)
     69     def add_point(self, x, y):
     70         index, x = x
---> 71         self.learners[index].add_point(x, y)
     72 
     73     def loss(self, real=True):

/opt/conda/lib/python3.6/site-packages/adaptive/learner/data_saver.py in add_point(self, x, result)
     34 
     35     def add_point(self, x, result):
---> 36         y = self.arg_picker(result)
     37         self.extra_data[x] = result
     38         self.learner.add_point(x, y)

TypeError: 'NoneType' object is not subscriptable