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

Closes #112 (closed)

adaptive/learner/learnerND.py

@@ -15,6 +15,15 @@ from .triangulation import (Triangulation, point_in_simplex,
                             circumsphere, simplex_volume_in_embedding)
 from ..utils import restore, cache_latest
 
+try:
+    import plotly.graph_objs
+    import plotly.figure_factory
+    import plotly.offline
+
+    with_plotly = True
+except ModuleNotFoundError:
+    with_plotly = False
+
 
 def volume(simplex, ys=None):
     # Notice the parameter ys is there so you can use this volume method as
@@ -576,3 +585,123 @@ class LearnerND(BaseLearner):
             return im.opts(style=dict(cmap='viridis'))
         else:
             raise ValueError("Only 1 or 2-dimensional plots can be generated.")
+
+    def get_range(self):
+        r_min = min(self.data[v] for v in self.tri.vertices)
+        r_max = max(self.data[v] for v in self.tri.vertices)
+        return r_min, r_max
+
+    def _get_isosurface(self, level=0.0):
+        if self.ndim != 3 or self.vdim != 1:
+            raise Exception('isosurface is only supported for 3d input and 1d output')
+
+        vertices = []            # index -> (x,y,z)
+        faces = []               # 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) # assuming they are numpy arrays
+            
+            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:
+            plane = []
+            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:
+                        if np.allclose(vertices[vi], vertices[pi]):
+                            should_add = False
+                    if should_add:
+                        plane.append(vi)
+                    
+            if len(plane) == 3:
+                faces.append(plane)
+            elif len(plane) == 4:
+                faces.append(plane[:3])
+                faces.append(plane[1:])
+
+        if len(faces) == 0:
+            r_min, r_max = self.get_range()
+
+            raise ValueError(f"""Could not draw isosurface for level={level}, as 
+            this value is not inside the function range. Please choose a level 
+            strictly inside interval ({r_min}, {r_max})""")
+
+        return vertices, faces
+
+    def plot_isosurface(self, level=0.0, hull_opacity=0.2):
+        """Plots the linearly interpolated iso-surface of the function, based on 
+        the currently evaluated points. This is the 3d analog of an iso-line.
+
+        Parameters
+        ----------
+        level : number
+            the function value which you are interested in. Defaults to 0.0.
+        hull_opacity : number
+            the opacity of the hull of the domain. Defaults to 0.2
+        """
+        if not with_plotly:
+            raise Exception('plot_isosurface requires plotly to be installed')
+
+        vertices, faces = self._get_isosurface(level)
+        x,y,z = zip(*vertices)
+
+        fig = plotly.figure_factory.create_trisurf(x=x,
+                                y=y, 
+                                z=z, 
+                                plot_edges=False,
+                                simplices=faces,
+                                title="Isosurface")
+
+        if hull_opacity < 1e-3:
+            # do not compute the hull_mesh
+            return plotly.offline.iplot(fig)
+
+        hull_mesh = self._get_hull_mesh(opacity=hull_opacity)
+        return plotly.offline.iplot([fig.data[0], hull_mesh])
+
+    def _get_hull_mesh(self, opacity=0.2):
+        hull = scipy.spatial.ConvexHull(self._bounds_points)
+
+        # find the colors of each plane, giving triangles which are coplanar the 
+        # same color, such that a square face has the same color.
+        color_dict = {}
+        def get_plane_color(simplex):
+            simplex= tuple(simplex)
+            # if the volume of the two triangles combined is zero then they belong to the same plane
+            for simplex_key, color in color_dict.items():
+                points = np.array([hull.points[i] for i in np.unique(simplex_key + simplex)])
+                if np.linalg.matrix_rank(points[1:] - points[0]) < 3:
+                    return color
+                if scipy.spatial.ConvexHull(points).volume < 1e-5:
+                    return color
+            color_dict[simplex] = (random.randint(0,255),random.randint(0,255),random.randint(0,255))
+            return color_dict[simplex]
+
+        colors = [get_plane_color(simplex) for simplex in hull.simplices]
+
+        x,y,z = zip(*self._bounds_points)
+        i,j,k = hull.simplices.T
+        return plotly.graph_objs.Mesh3d(x=x, y=y, z=z, i=i, j=j, k=k, facecolor=(colors), opacity=opacity)
+        
\ No newline at end of file