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

Closes #112 (closed)

adaptive/learner/learnerND.py

@@ -10,20 +10,11 @@ import scipy.spatial
 
 from .base_learner import BaseLearner
 
-from ..notebook_integration import ensure_holoviews
+from ..notebook_integration import ensure_holoviews, ensure_plotly
 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
@@ -472,6 +463,10 @@ class LearnerND(BaseLearner):
         self._subtriangulations = dict()
         self._pending_to_simplex = dict()
 
+    ##########################
+    # Plotting related stuff #
+    ##########################
+
     def plot(self, n=None, tri_alpha=0):
         """Plot the function we want to learn, only works in 2D.
 
@@ -586,55 +581,51 @@ class LearnerND(BaseLearner):
         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')
+            raise Exception('Isosurface plotting is only supported'
+                            ' for a 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
 
-        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)]
 
-        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
+            # 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]
+            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_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
+            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):
+            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)
+                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:
@@ -642,66 +633,172 @@ class LearnerND(BaseLearner):
                 faces.append(plane[1:])
 
         if len(faces) == 0:
-            r_min, r_max = self.get_range()
+            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 a level 
-            strictly inside interval ({r_min}, {r_max})""")
+            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, 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
+
+    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
+        learn, only works in 2D.
+
+        Parameters
+        ----------
+        level : float
+            the value of the function at which you would like to see the isoline
+        n : int
+            the number of boxes in the interpolation grid along each axis
+        """
+        hv = ensure_holoviews()
+        if n == -1:
+            plot = hv.Path([])
+        else:
+            plot = self.plot(n=n, tri_alpha=tri_alpha)
+        
+        if isinstance(level, Iterable):
+            for l in level:
+                plot = plot * self.plot_isoline(level=l, n=-1)
+            return plot
+
+
+        vertices, lines = self._get_isoline(level)
+        paths = [[vertices[i], vertices[j]] for i,j in lines]
+        contour = hv.Path(paths)
+
+        contour_opts = dict(color='black')
+        contour = contour.opts(style=contour_opts)
+        return plot * contour
 
     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.
+        """Plots the linearly interpolated isosurface of the function,
+        based on the currently evaluated points. This is the 3D analog
+        of an isoline.
 
         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
+        level : float, default 0.0
+            the function value which you are interested in.
+        hull_opacity : float, default 0.0
+            the opacity of the hull of the domain.
+
+        Returns
+        -------
+        plot : plotly.offline.iplot object
+            The plot object of the isosurface.
         """
-        if not with_plotly:
-            raise Exception('plot_isosurface requires plotly to be installed')
+        plotly = ensure_plotly()
 
         vertices, faces = self._get_isosurface(level)
-        x,y,z = zip(*vertices)
+        x, y, z = zip(*vertices)
 
-        fig = plotly.figure_factory.create_trisurf(x=x,
-                                y=y, 
-                                z=z, 
-                                plot_edges=False,
-                                simplices=faces,
-                                title="Isosurface")
+        fig = plotly.figure_factory.create_trisurf(
+            x=x, y=y, z=z, plot_edges=False,
+            simplices=faces, title="Isosurface")
+        isosurface = fig.data[0]
+        isosurface.update(lighting=dict(ambient=1, diffuse=1, roughness=1, specular=0, fresnel=0))
 
         if hull_opacity < 1e-3:
-            # do not compute the hull_mesh
+            # 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])
+        return plotly.offline.iplot([isosurface, hull_mesh])
 
     def _get_hull_mesh(self, opacity=0.2):
+        plotly = ensure_plotly()
         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.
+        # 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
+
+        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)])
+                points = [hull.points[i] for i in set(simplex_key + simplex)]
+                points = np.array(points)
                 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))
+            color_dict[simplex] = tuple(random.randint(0, 255)
+                                        for _ in range(3))
             return color_dict[simplex]
 
-        colors = [get_plane_color(simplex) for simplex in hull.simplices]
+        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
+        x, y, z = zip(*self._bounds_points)
+        i, j, k = hull.simplices.T
+        lighting = dict(ambient=1, diffuse=1, roughness=1, specular=0, fresnel=0)
+        return plotly.graph_objs.Mesh3d(x=x, y=y, z=z, i=i, j=j, k=k,
+                                        facecolor=colors, opacity=opacity,
+                                        lighting=lighting)