Repository: PMEAL/OpenPNM

Issue:
Should plotting functions be moved to a different place?
We have tended to use ``topotools`` as a bit of a dumping ground for functions that pertain to "network stuff".  We have 4-5 plotting functions in there.  I'm writing the plotting tutorial and wondering if these could be in a better spot?  Not sure where, so any thoughts?

Discussion:

Patch:
diff --git a/openpnm/__init__.py b/openpnm/__init__.py
index e6b9a0ef9..11ec6d995 100644
--- a/openpnm/__init__.py
+++ b/openpnm/__init__.py
@@ -20,6 +20,7 @@ from . import solvers
 from . import integrators
 from . import io
 from . import contrib
+from . import visualization
 
 from .utils import Workspace, Project
 
diff --git a/openpnm/contrib/__init__.py b/openpnm/contrib/__init__.py
index 4144f1433..994381c77 100644
--- a/openpnm/contrib/__init__.py
+++ b/openpnm/contrib/__init__.py
@@ -11,4 +11,3 @@ removed.
 
 from ._transient_multiphysics import *
 from ._multiphase import MultiPhase
-from ._plot_network import *
diff --git a/openpnm/contrib/_plot_network.py b/openpnm/contrib/_plot_network.py
deleted file mode 100644
index c40595dc4..000000000
--- a/openpnm/contrib/_plot_network.py
+++ /dev/null
@@ -1,109 +0,0 @@
-import numpy as np
-import openpnm as op
-from matplotlib import cm
-
-
-__all__ = [
-    'plot_3d',
-]
-
-
-def create_pore_colors_from_array(a, cmap='viridis'):
-    colormap = cm.get_cmap(cmap)
-    return colormap(a/a.max())
-
-
-def create_throat_colors_from_array(a, cmap='viridis'):
-    colormap = cm.get_cmap(cmap)
-    return np.repeat(colormap(a/a.max()), 2, axis=0)
-
-
-def plot_3d(
-    network,
-    pore_color=None,
-    pore_size=None,
-    throat_color=None,
-    throat_size=None,
-    bgcolor='grey',
-):
-    r"""
-
-    Parameters
-    ----------
-    network
-    """
-    try:
-        from vispy import scene
-    except ModuleNotFoundError:
-        raise Exception("vispy must be installed to use this function")
-    canvas = scene.SceneCanvas(keys='interactive', show=True, bgcolor=bgcolor)
-    view = canvas.central_widget.add_view()
-    view.camera = 'turntable'
-    view.camera.fov = 30
-    view.camera.distance = 3*np.max(network['pore.coords'])
-
-    if pore_color is None:
-        pore_color = create_pore_colors_from_array(network['pore.diameter'],
-                                                   cmap='viridis')
-    else:
-        pore_color = create_pore_colors_from_array(pore_color,
-                                                   cmap='viridis')
-    if throat_color is None:
-        throat_color = create_throat_colors_from_array(network['throat.diameter'],
-                                                       cmap='viridis')
-    else:
-        throat_color = create_throat_colors_from_array(throat_color,
-                                                       cmap='viridis')
-    if pore_size is None:
-        pore_size = network['pore.diameter']
-    if throat_size is None:
-        throat_size = 2
-    else:
-        throat_size = np.max(throat_size)  # Arrays not supported here
-    # plot spheres
-    vis = scene.visuals.Markers(
-        pos=network['pore.coords'],
-        size=pore_size,
-        antialias=0,
-        face_color=pore_color,
-        edge_width=0,
-        scaling=True,
-        spherical=True,
-    )
-    vis.parent = view.scene
-    # plot axis
-    # vispy.scene.visuals.XYZAxis(parent=view.scene)
-    # set camera center
-    view.camera.center = np.array((network['pore.coords'][:, 0].max()/2,
-                                   network['pore.coords'][:, 1].max()/2,
-                                   network['pore.coords'][:, 2].max()/2))
-    # data preparation
-    lines = np.zeros((len(network['throat.conns']), 2, 3))
-    for i in range(len(network['throat.conns'])):
-        pair = network['throat.conns'][i]
-        line = np.array([[network['pore.coords'][pair[0]],
-                          network['pore.coords'][pair[1]]]])
-        lines[i, :, :] = line
-    # plot throats
-    vis2 = scene.visuals.Line(lines,
-                              width=throat_size,
-                              color=throat_color,
-                              connect='segments',
-                              antialias=True,)
-    vis2.parent = view.scene
-
-
-if __name__ == '__main__':
-    pn = op.network.Cubic(shape=[30, 30, 30], spacing=1)
-    pn['pore.diameter'] = np.random.rand(pn.Np)
-    print(f'Total number of pores: {pn.Np}')
-    print(f'Total number of throats: {pn.Nt}')
-    P12 = pn['throat.conns']
-    D12 = pn['pore.diameter'][P12]
-    Dt = np.amin(D12, axis=1)
-    pn['throat.diameter'] = Dt
-    plot_3d(network=pn,
-            pore_color=pn['pore.diameter'],
-            throat_color=pn['throat.diameter'],
-            throat_size=10*pn['throat.diameter'],
-            bgcolor='grey')
diff --git a/openpnm/topotools/__init__.py b/openpnm/topotools/__init__.py
index 4b67f0d6b..2836bb9e4 100644
--- a/openpnm/topotools/__init__.py
+++ b/openpnm/topotools/__init__.py
@@ -10,4 +10,3 @@ network topology.
 from ._topotools import *
 from ._perctools import *
 from ._graphtools import *
-from ._plottools import *
diff --git a/openpnm/visualization/__init__.py b/openpnm/visualization/__init__.py
new file mode 100644
index 000000000..3905ff704
--- /dev/null
+++ b/openpnm/visualization/__init__.py
@@ -0,0 +1,1 @@
+from ._plottools import *
diff --git a/openpnm/topotools/_plottools.py b/openpnm/visualization/_plottools.py
similarity index 84%
rename from openpnm/topotools/_plottools.py
rename to openpnm/visualization/_plottools.py
index 5762bfd49..6543d2f72 100644
--- a/openpnm/topotools/_plottools.py
+++ b/openpnm/visualization/_plottools.py
@@ -1,6 +1,7 @@
 import numpy as np
 import openpnm as op
 from tqdm import tqdm
+from matplotlib.pyplot import cm
 
 
 __all__ = [
@@ -8,7 +9,8 @@ __all__ = [
     'plot_coordinates',
     'plot_networkx',
     'plot_tutorial',
-    'plot_network_jupyter',
+    'plot_notebook',
+    'plot_vispy',
     'generate_voxel_image',
 ]
 
@@ -93,13 +95,13 @@ def plot_connections(network,
     >>> pn.add_boundary_pores()
     >>> Ts = pn.throats('*boundary', mode='not')  # find internal throats
     >>> fig, ax = plt.subplots()  # create empty figure
-    >>> _ = op.topotools.plot_connections(network=pn,
-    ...                                   throats=Ts)  # plot internal throats
+    >>> _ = op.visualization.plot_connections(network=pn,
+    ...                                       throats=Ts)  # plot internal throats
     >>> Ts = pn.throats('*boundary')  # find boundary throats
-    >>> _ = op.topotools.plot_connections(network=pn,
-    ...                                   throats=Ts,
-    ...                                   ax=ax,
-    ...                                   color='r')  # plot boundary throats in red
+    >>> _ = op.visualization.plot_connections(network=pn,
+    ...                                       throats=Ts,
+    ...                                       ax=ax,
+    ...                                       color='r')  # plot boundary throats in red
 
     """
     import matplotlib.pyplot as plt
@@ -242,15 +244,15 @@ def plot_coordinates(network,
     >>> pn.add_boundary_pores()
     >>> Ps = pn.pores('internal')  # find internal pores
     >>> fig, ax = plt.subplots()  # create empty figure
-    >>> _ = op.topotools.plot_coordinates(network=pn,
-    ...                                   pores=Ps,
-    ...                                   color='b',
-    ...                                   ax=ax)  # plot internal pores
+    >>> _ = op.visualization.plot_coordinates(network=pn,
+    ...                                       pores=Ps,
+    ...                                       color='b',
+    ...                                       ax=ax)  # plot internal pores
     >>> Ps = pn.pores('*boundary')  # find boundary pores
-    >>> _ = op.topotools.plot_coordinates(network=pn,
-    ...                                   pores=Ps,
-    ...                                   color='r',
-    ...                                   ax=ax)  # plot boundary pores in red
+    >>> _ = op.visualization.plot_coordinates(network=pn,
+    ...                                       pores=Ps,
+    ...                                       color='r',
+    ...                                       ax=ax)  # plot boundary pores in red
 
     """
     import matplotlib.pyplot as plt
@@ -298,11 +300,21 @@ def plot_coordinates(network,
         markersize = size_by / size_by.max() * markersize
 
     if ThreeD:
-        sc = ax.scatter(X, Y, Z, c=color, s=markersize, marker=marker, alpha=alpha, **kwargs)
+        sc = ax.scatter(X, Y, Z,
+                        c=color,
+                        s=markersize,
+                        marker=marker,
+                        alpha=alpha,
+                        **kwargs)
         _scale_axes(ax=ax, X=Xl, Y=Yl, Z=Zl)
     else:
         _X, _Y = np.column_stack((X, Y, Z))[:, dim].T
-        sc = ax.scatter(_X, _Y, c=color, s=markersize, marker=marker, alpha=alpha, **kwargs)
+        sc = ax.scatter(_X, _Y,
+                        c=color,
+                        s=markersize,
+                        marker=marker,
+                        alpha=alpha,
+                        **kwargs)
         _scale_axes(ax=ax, X=Xl, Y=Yl, Z=np.zeros_like(Yl))
 
     _label_axes(ax=ax, X=Xl, Y=Yl, Z=Zl)
@@ -504,13 +516,13 @@ def plot_tutorial(network,
     return gplot
 
 
-def plot_network_jupyter(network,
-                         node_color=0,
-                         edge_color=0,
-                         node_size=1,
-                         node_scale=20,
-                         edge_scale=5,
-                         colormap='viridis'):
+def plot_notebook(network,
+                  node_color=0,
+                  edge_color=0,
+                  node_size=1,
+                  node_scale=20,
+                  edge_scale=5,
+                  colormap='viridis'):
     r"""
     Visualize a network in 3D using Plotly.
 
@@ -567,8 +579,9 @@ def plot_network_jupyter(network,
     node_color = np.ones(network.Np)*node_color
     edge_color = np.ones(network.Nt)*edge_color
 
-    node_labels = [str(i)+ ': ' + str(x) for i, x in enumerate(zip(node_size, node_color))]
-    edge_labels = [str(i)+ ': ' + str(x) for i, x in enumerate(edge_color)]
+    node_labels = [str(i) + ': ' + str(x) for i, x in
+                   enumerate(zip(node_size, node_color))]
+    edge_labels = [str(i) + ': ' + str(x) for i, x in enumerate(edge_color)]
 
     # Create edges and nodes coordinates
     N = network.Nt*3
@@ -731,7 +744,7 @@ def _generate_voxel_image(network, pore_shape, throat_shape, max_dim=200):
 def generate_voxel_image(network, pore_shape="sphere", throat_shape="cylinder",
                          max_dim=None, rtol=0.1):
     r"""
-    Generate a voxel image from an Network
+    Generate a voxel image from a Network
 
     Parameters
     ----------
@@ -778,3 +791,88 @@ def generate_voxel_image(network, pore_shape="sphere", throat_shape="cylinder",
         eps_old = eps
         max_dim = int(max_dim * 1.25)
     return im
+
+
+def create_pore_colors_from_array(a, cmap='viridis'):
+    colormap = cm.get_cmap(cmap)
+    return colormap(a/a.max())
+
+
+def create_throat_colors_from_array(a, cmap='viridis'):
+    colormap = cm.get_cmap(cmap)
+    return np.repeat(colormap(a/a.max()), 2, axis=0)
+
+
+def plot_vispy(
+    network,
+    pore_color=None,
+    pore_size=None,
+    throat_color=None,
+    throat_size=None,
+    bgcolor='grey',
+):
+    r"""
+
+    Parameters
+    ----------
+    network
+    """
+    try:
+        from vispy import scene
+    except ModuleNotFoundError:
+        raise Exception("vispy must be installed to use this function")
+    canvas = scene.SceneCanvas(keys='interactive', show=True, bgcolor=bgcolor)
+    view = canvas.central_widget.add_view()
+    view.camera = 'turntable'
+    view.camera.fov = 30
+    view.camera.distance = 3*np.max(network['pore.coords'])
+
+    if pore_color is None:
+        pore_color = create_pore_colors_from_array(network['pore.diameter'],
+                                                   cmap='viridis')
+    else:
+        pore_color = create_pore_colors_from_array(pore_color,
+                                                   cmap='viridis')
+    if throat_color is None:
+        throat_color = create_throat_colors_from_array(network['throat.diameter'],
+                                                       cmap='viridis')
+    else:
+        throat_color = create_throat_colors_from_array(throat_color,
+                                                       cmap='viridis')
+    if pore_size is None:
+        pore_size = network['pore.diameter']
+    if throat_size is None:
+        throat_size = 2
+    else:
+        throat_size = np.max(throat_size)  # Arrays not supported here
+    # plot spheres
+    vis = scene.visuals.Markers(
+        pos=network['pore.coords'],
+        size=pore_size,
+        antialias=0,
+        face_color=pore_color,
+        edge_width=0,
+        scaling=True,
+        spherical=True,
+    )
+    vis.parent = view.scene
+    # plot axis
+    # vispy.scene.visuals.XYZAxis(parent=view.scene)
+    # set camera center
+    view.camera.center = np.array((network['pore.coords'][:, 0].max()/2,
+                                   network['pore.coords'][:, 1].max()/2,
+                                   network['pore.coords'][:, 2].max()/2))
+    # data preparation
+    lines = np.zeros((len(network['throat.conns']), 2, 3))
+    for i in range(len(network['throat.conns'])):
+        pair = network['throat.conns'][i]
+        line = np.array([[network['pore.coords'][pair[0]],
+                          network['pore.coords'][pair[1]]]])
+        lines[i, :, :] = line
+    # plot throats
+    vis2 = scene.visuals.Line(lines,
+                              width=throat_size,
+                              color=throat_color,
+                              connect='segments',
+                              antialias=True,)
+    vis2.parent = view.scene


Test patch:
diff --git a/tests/unit/topotools/PlotToolsTest.py b/tests/unit/visualization/PlotToolsTest.py
similarity index 87%
rename from tests/unit/topotools/PlotToolsTest.py
rename to tests/unit/visualization/PlotToolsTest.py
index 1d7ceea87..9f29c6bff 100644
--- a/tests/unit/topotools/PlotToolsTest.py
+++ b/tests/unit/visualization/PlotToolsTest.py
@@ -28,7 +28,7 @@ class PlotToolsTest:
             dims = op.topotools.dimensionality(pn)
             x, y = pn["pore.coords"].T[dims]
             fig, ax = plt.subplots()
-            m = op.topotools.plot_networkx(pn, ax=ax)
+            m = op.visualization.plot_networkx(pn, ax=ax)
             x_plot, y_plot = np.array(m.get_offsets()).T
             np.testing.assert_allclose(x_plot, x)
             np.testing.assert_allclose(y_plot, y)
@@ -42,7 +42,7 @@ class PlotToolsTest:
             pn = op.network.Cubic(shape=shape)
             x, y = pn["pore.coords"].T[op.topotools.dimensionality(pn)]
             fig, ax = plt.subplots()
-            m = op.topotools.plot_networkx(pn, ax=ax)
+            m = op.visualization.plot_networkx(pn, ax=ax)
             x_plot, y_plot = np.array(m.get_offsets()).T
             np.testing.assert_allclose(x_plot, x)
             np.testing.assert_allclose(y_plot, y)
@@ -54,7 +54,7 @@ class PlotToolsTest:
             pn = op.network.Cubic(shape=shape)
             x, = pn["pore.coords"].T[op.topotools.dimensionality(pn)]
             fig, ax = plt.subplots()
-            m = op.topotools.plot_networkx(pn, ax=ax)
+            m = op.visualization.plot_networkx(pn, ax=ax)
             x_plot, y_plot = np.array(m.get_offsets()).T
             np.testing.assert_allclose(x_plot, x)
             plt.close()
@@ -62,14 +62,14 @@ class PlotToolsTest:
     def test_plot_networkx_3d(self):
         pn = op.network.Cubic(shape=[5, 8, 3])
         with pytest.raises(Exception):
-            op.topotools.plot_networkx(pn)
+            op.visualization.plot_networkx(pn)
 
     def test_generate_voxel_image(self):
         pn = op.network.Cubic(shape=[5, 5, 1])
         pn.add_model_collection(
             op.models.collections.geometry.spheres_and_cylinders)
         pn.regenerate_models()
-        im = op.topotools.generate_voxel_image(network=pn,
+        im = op.visualization.generate_voxel_image(network=pn,
                                                pore_shape='sphere',
                                                throat_shape='cylinder',
                                                max_dim=500)
@@ -82,8 +82,8 @@ class PlotToolsTest:
             op.models.collections.geometry.spheres_and_cylinders)
         pn.regenerate_models()
         Ts = np.array([0, 4, 6, 18])
-        im = op.topotools.plot_connections(pn, throats=Ts,
-                                           color_by=pn['throat.diameter'])
+        im = op.visualization.plot_connections(pn, throats=Ts,
+                                               color_by=pn['throat.diameter'])
         colors_im = im.get_color()
         color_by = pn['throat.diameter'][Ts]
         color_calc = cm.get_cmap(name='jet')(color_by / color_by.max())
@@ -97,8 +97,8 @@ class PlotToolsTest:
             op.models.collections.geometry.spheres_and_cylinders)
         pn.regenerate_models()
         Ps = np.array([0, 4, 6, 18])
-        im = op.topotools.plot_coordinates(pn, pores=Ps,
-                                           color_by=pn['pore.diameter'])
+        im = op.visualization.plot_coordinates(pn, pores=Ps,
+                                               color_by=pn['pore.diameter'])
         colors_im = im.get_edgecolors()
         color_by = pn['pore.diameter'][Ps]
         color_calc = cm.get_cmap(name='jet')(color_by / color_by.max())