Repository: pmeal/openpnm
Issue: color_by argument in plot_network_connections/coordinates does not respect throats/pores argument

I you only plot X throats, but pass a full length array for ``color_by``, then the function will assign the first X entries in ``color_by`` to the throats.  It should either throw an error that the two are the wrong length, or use the ``throats/pores`` list to index into the color array.

Changed files:
openpnm/topotools/_plottools.py, tests/unit/topotools/PlotToolsTest.py

diff --git a/openpnm/topotools/_plottools.py b/openpnm/topotools/_plottools.py
index 99aa5bd8a..2dc4f1d43 100644
--- a/openpnm/topotools/_plottools.py
+++ b/openpnm/topotools/_plottools.py
@@ -142,6 +142,7 @@ def plot_connections(network,
     color = mcolors.to_rgb(color) + tuple([alpha])
     # Override colors with color_by if given
     if color_by is not None:
+        color_by = color_by[Ts]
         color = cm.get_cmap(name=cmap)(color_by / color_by.max())
         color[:, 3] = alpha
     if size_by is not None:
@@ -289,6 +290,7 @@ def plot_coordinates(network,
     if 's' in kwargs.keys():
         markersize = kwargs.pop('s')
     if color_by is not None:
+        color_by = color_by[Ps]
         color = cm.get_cmap(name=cmap)(color_by / color_by.max())
     if size_by is not None:
         markersize = size_by / size_by.max() * markersize
diff --git a/tests/unit/topotools/PlotToolsTest.py b/tests/unit/topotools/PlotToolsTest.py
index 415a048e0..21290baea 100644
--- a/tests/unit/topotools/PlotToolsTest.py
+++ b/tests/unit/topotools/PlotToolsTest.py
@@ -2,7 +2,9 @@ import pytest
 import numpy as np
 import openpnm as op
 import matplotlib.pyplot as plt
+from matplotlib import cm
 from openpnm import topotools
+from numpy.testing import assert_allclose
 
 
 class PlotToolsTest:
@@ -71,6 +73,34 @@ class PlotToolsTest:
                                                max_dim=500)
         assert im.shape[0] == 500
 
+    def test_plot_connections_color_by(self):
+        pn = op.network.Cubic(shape=[5, 5, 1])
+        np.random.seed(10)
+        geo = op.geometry.SpheresAndCylinders(network=pn,
+                                              pores=pn.Ps, throats=pn.Ts)
+        Ts = np.array([0, 4, 6, 18])
+        im = op.topotools.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())
+        color_calc[:, 3] = 1.0
+        assert_allclose(color_calc, colors_im, rtol=1e-5)
+
+    def test_plot_coordinates_color_by(self):
+        pn = op.network.Cubic(shape=[5, 5, 1])
+        np.random.seed(10)
+        geo = op.geometry.SpheresAndCylinders(network=pn,
+                                              pores=pn.Ps, throats=pn.Ts)
+        Ps = np.array([0, 4, 6, 18])
+        im = op.topotools.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())
+        color_calc[:, 3] = 1.0
+        assert_allclose(color_calc, colors_im, rtol=1e-5)
+
 
 if __name__ == '__main__':