Resolve "(LearnerND) add iso-surface plot feature"

Closes #112 (closed)

adaptive/learner/learnerND.py

@@ -658,13 +658,22 @@ class LearnerND(BaseLearner):
     def _set_data(self, data):
         self.tell_many(*zip(*data.items()))
 
-    def _get_isosurface(self, level=0.0):
-        if self.ndim != 3 or self.vdim != 1:
-            raise Exception('Isosurface plotting is only supported'
-                            ' for a 3D input and 1D output')
+    def _get_iso(self, level=0.0, which='surface'):
+        if which == 'surface':
+            if self.ndim != 3 or self.vdim != 1:
+                raise Exception('Isosurface plotting is only supported'
+                                ' for a 3D input and 1D output')
+            get_surface = True
+            get_line = False
+        elif which == 'line':
+            if self.ndim != 2 or self.vdim != 1:
+                raise Exception('Isoline plotting is only supported'
+                                ' for a 2D input and 1D output')
+            get_surface = False
+            get_line = True
 
         vertices = []  # index -> (x,y,z)
-        faces = []  # tuple of indices of the corner points
+        faces_or_lines = []  # tuple of indices of the corner points
         from_line_to_vertex = {}  # the interpolated vertex (index) between two known points
 
         def _get_vertex_index(a, b):
@@ -690,26 +699,28 @@ class LearnerND(BaseLearner):
             return new_index
 
         for simplex in self.tri.simplices:
-            plane = []
+            plane_or_line = []
             for a, b in itertools.combinations(simplex, 2):
                 va = self.data[self.tri.vertices[a]]
                 vb = self.data[self.tri.vertices[b]]
                 if min(va, vb) < level <= max(va, vb):
                     vi = _get_vertex_index(a, b)
                     should_add = True
-                    for pi in plane:
+                    for pi in plane_or_line:
                         if np.allclose(vertices[vi], vertices[pi]):
                             should_add = False
                     if should_add:
-                        plane.append(vi)
+                        plane_or_line.append(vi)
 
-            if len(plane) == 3:
-                faces.append(plane)
-            elif len(plane) == 4:
-                faces.append(plane[:3])
-                faces.append(plane[1:])
+            if get_surface and len(plane_or_line) == 3:
+                faces_or_lines.append(plane_or_line)
+            elif get_surface and len(plane_or_line) == 4:
+                faces_or_lines.append(plane_or_line[:3])
+                faces_or_lines.append(plane_or_line[1:])
+            elif get_line and len(plane_or_line) == 2:
+                faces_or_lines.append(plane_or_line)
 
-        if len(faces) == 0:
+        if len(faces_or_lines) == 0:
             r_min = min(self.data[v] for v in self.tri.vertices)
             r_max = max(self.data[v] for v in self.tri.vertices)
 
@@ -719,68 +730,7 @@ class LearnerND(BaseLearner):
                 f" a level strictly inside interval ({r_min}, {r_max})"
             )
 
-        return vertices, faces
-
-    def _get_isoline(self, level=0.0):
-        # Very similar to _get_isosurface, maybe merge the two functions
-        if self.ndim != 2 or self.vdim != 1:
-            raise Exception('Isoline plotting is only supported'
-                            ' for a 2D input and 1D output')
-
-        vertices = []  # index -> (x,y,z)
-        lines = []  # tuple of indices of the corner points
-        from_line_to_vertex = {}  # the interpolated vertex (index) between two known points
-
-        def _get_vertex_index(a, b):
-            if (a, b) in from_line_to_vertex:
-                return from_line_to_vertex[(a, b)]
-
-            # Otherwise compute it and cache the result.
-            vertex_a = self.tri.vertices[a]
-            vertex_b = self.tri.vertices[b]
-            value_a = self.data[vertex_a]
-            value_b = self.data[vertex_b]
-
-            da = abs(value_a - level)
-            db = abs(value_b - level)
-            dab = da + db
-
-            new_pt = (db / dab * np.array(vertex_a)
-                      + da / dab * np.array(vertex_b))
-
-            new_index = len(vertices)
-            vertices.append(new_pt)
-            from_line_to_vertex[(a, b)] = new_index
-            return new_index
-
-        for simplex in self.tri.simplices:
-            line = []
-            for a, b in itertools.combinations(simplex, 2):
-                va = self.data[self.tri.vertices[a]]
-                vb = self.data[self.tri.vertices[b]]
-                if min(va, vb) < level <= max(va, vb):
-                    vi = _get_vertex_index(a, b)
-                    should_add = True
-                    for pi in line:
-                        if np.allclose(vertices[vi], vertices[pi]):
-                            should_add = False
-                    if should_add:
-                        line.append(vi)
-
-            if len(line) == 2:
-                lines.append(line)
-
-        if len(lines) == 0:
-            r_min = min(self.data[v] for v in self.tri.vertices)
-            r_max = max(self.data[v] for v in self.tri.vertices)
-
-            raise ValueError(
-                f"Could not draw isosurface for level={level}, as"
-                " this value is not inside the function range. Please choose"
-                f" a level strictly inside interval ({r_min}, {r_max})"
-            )
-
-        return vertices, lines
+        return vertices, faces_or_lines
 
     def plot_isoline(self, level=0.0, n=None, tri_alpha=0):
         """Plot the isoline at a specific level of the function we want to
@@ -804,8 +754,7 @@ class LearnerND(BaseLearner):
                 plot = plot * self.plot_isoline(level=l, n=-1)
             return plot
 
-
-        vertices, lines = self._get_isoline(level)
+        vertices, lines = self.self._get_iso(level, which='line')
         paths = [[vertices[i], vertices[j]] for i, j in lines]
         contour = hv.Path(paths)
 
@@ -832,7 +781,7 @@ class LearnerND(BaseLearner):
         """
         plotly = ensure_plotly()
 
-        vertices, faces = self._get_isosurface(level)
+        vertices, faces = self._get_iso(level, which='surface')
         x, y, z = zip(*vertices)
 
         fig = plotly.figure_factory.create_trisurf(