test flat bands in the LearnerND

Author: basnijholt

adaptive/tests/test_learnernd.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 
+import numpy as np
 import scipy.spatial
 
 from adaptive.learner import LearnerND
@@ -38,3 +39,14 @@ def test_interior_vs_bbox_gives_same_result():
     simple(learner, goal=lambda l: l.loss() < 0.1)
 
     assert learner.data == control.data
+
+
+def test_vector_return_with_a_flat_layer():
+    f = generate_random_parametrization(ring_of_fire)
+    g = generate_random_parametrization(ring_of_fire)
+    h1 = lambda xy: np.array([f(xy), g(xy)])
+    h2 = lambda xy: np.array([f(xy), 0*g(xy)])
+    h3 = lambda xy: np.array([0*f(xy), g(xy)])
+    for function in [h1, h2, h3]:
+        learner = LearnerND(function, bounds=[(-1, 1), (-1, 1)])
+        simple(learner, goal=lambda l: l.loss() < 0.1)