WIP: allow the loss to be a tuple in the Learner1D

Author: basnijholt

adaptive/learner/learner1D.py

@@ -69,6 +69,22 @@ def _wrapped(loss_per_interval):
     return _wrapped
 
 
+def loss_returns(return_type, return_length):
+    def _wrapped(loss_per_interval):
+        loss_per_interval.return_type = return_type
+        loss_per_interval.return_length = return_length
+        return loss_per_interval
+    return _wrapped
+
+
+def inf_format(return_type, return_len=None):
+    is_iterable = hasattr(return_type, '__iter__')
+    if is_iterable:
+        return return_type(return_len * [np.inf])
+    else:
+        return return_type(np.inf)
+
+
 @uses_nth_neighbors(0)
 def uniform_loss(xs, ys):
     """Loss function that samples the domain uniformly.
@@ -287,7 +303,8 @@ def npoints(self):
     def loss(self, real=True):
         losses = self.losses if real else self.losses_combined
         if not losses:
-            return np.inf
+            return inf_format(self.loss_per_interval.return_type,
+                              self.loss_per_interval.return_length)
         max_interval, max_loss = losses.peekitem(0)
         return max_loss
 
@@ -660,8 +677,14 @@ def _set_data(self, data):
 
 def loss_manager(x_scale):
     def sort_key(ival, loss):
-        loss, ival = finite_loss(ival, loss, x_scale)
-        return -loss, ival
+        if isinstance(loss, Iterable):
+            loss, ival = zip(*[finite_loss(ival, l, x_scale) for l in loss])
+            loss = tuple(-x for x in loss)
+            ival = ival[0]
+        else:
+            loss, ival = finite_loss(ival, loss, x_scale)
+            loss = -loss
+        return loss, ival
     sorted_dict = sortedcollections.ItemSortedDict(sort_key)
     return sorted_dict