Merge pull request yhat#339 from matthias-k/geom_pointrange

Geom pointrange

Author: glamp

ggplot/geoms/__init__.py

@@ -10,9 +10,11 @@
 from .geom_hline import geom_hline
 from .geom_jitter import geom_jitter
 from .geom_line import geom_line
+from .geom_linerange import geom_linerange
 from .geom_now_its_art import geom_now_its_art
 from .geom_path import geom_path
 from .geom_point import geom_point
+from .geom_pointrange import geom_pointrange
 from .geom_rect import geom_rect
 from .geom_smooth import geom_smooth
 from .geom_step import geom_step
@@ -27,8 +29,8 @@
 
 __facet__ = ['facet_grid', 'facet_wrap']
 __geoms__ = ['geom_abline', 'geom_area', 'geom_bar', 'geom_boxplot', 'geom_density',
-             'geom_histogram', 'geom_hline', 'geom_jitter', 'geom_line',
-             'geom_now_its_art', 'geom_path', 'geom_point', 'geom_rect',
+             'geom_histogram', 'geom_hline', 'geom_jitter', 'geom_line', 'geom_linerange',
+             'geom_now_its_art', 'geom_path', 'geom_point', 'geom_pointrange', 'geom_rect',
              'geom_step', 'geom_smooth', 'geom_text', 'geom_tile',
              'geom_vline']
 __components__ = ['ylab', 'xlab', 'ylim', 'xlim', 'labs', 'ggtitle']

ggplot/geoms/geom_linerange.py

@@ -0,0 +1,95 @@
+from __future__ import (absolute_import, division, print_function,
+                        unicode_literals)
+
+import sys
+
+from .geom import geom
+from ggplot.utils import is_categorical
+import numpy as np
+
+
+class geom_linerange(geom):
+    """Plot intervals represented by vertical lines
+
+    Parameters
+    ---------
+
+    x
+        x values of data
+    ymin
+        lower end of the interval for each x
+    ymax
+        upper end of the interval for each x
+    alpha : float
+        alpha value, defaults to 1
+    color : string
+        line color, defaults to 'black'
+    linetype : string
+        line type, defaults to 'solid'
+    size : string
+        width of the line, defaults to 2
+
+    Examples
+    --------
+
+    .. plot::
+        :include-source:
+
+        import numpy as np
+        import pandas as pd
+        from ggplot import *
+
+        np.random.seed(42)
+        x = np.linspace(0.5, 9.5, num=10)
+        y = np.random.randn(10)
+        ymin = y - np.random.uniform(0,1, size=10)
+        ymax = y + np.random.uniform(0,1, size=10)
+
+        data = pd.DataFrame({'x': x, 'ymin': ymin, 'ymax': ymax})
+
+        ggplot(aes(x='x', ymin='ymin', ymax='ymax'), data) \
+            + geom_linerange()
+
+    """
+    DEFAULT_AES = {'alpha': 1, 'color': 'black',
+                   'linetype': 'solid',
+                   'size': 2}
+    REQUIRED_AES = {'x', 'ymin', 'ymax'}
+    DEFAULT_PARAMS = {'stat': 'identity', 'position': 'identity', 'cmap': None}
+
+    _aes_renames = {'size': 'linewidth', 'linetype': 'linestyle'}
+    _units = {'alpha', 'color', 'linestyle'}
+
+    def __init__(self, *args, **kwargs):
+        super(geom_linerange, self).__init__(*args, **kwargs)
+        self._warning_printed = False
+
+    def _plot_unit(self, pinfo, ax):
+        # If x is categorical, calculate positions to plot
+        categorical = is_categorical(pinfo['x'])
+        if categorical:
+            x = pinfo.pop('x')
+            new_x = np.arange(len(x))
+            ax.set_xticks(new_x)
+            ax.set_xticklabels(x)
+            pinfo['x'] = new_x
+
+        if 'linewidth' in pinfo and isinstance(pinfo['linewidth'], list):
+            # ggplot also supports aes(size=...) but the current mathplotlib
+            # is not. See https://github.com/matplotlib/matplotlib/issues/2658
+            pinfo['linewidth'] = 4
+            if not self._warning_printed:
+                msg = "'geom_line()' currenty does not support the mapping of " +\
+                      "size ('aes(size=<var>'), using size=4 as a replacement.\n" +\
+                      "Use 'geom_line(size=x)' to set the size for the whole line.\n"
+                sys.stderr.write(msg)
+                self._warning_printed = True
+
+        x = pinfo.pop('x')
+        x = np.vstack([x, x])
+
+        ymin = pinfo.pop('ymin')
+        ymax = pinfo.pop('ymax')
+        y = np.vstack([ymin, ymax])
+
+        ax.plot(x, y, **pinfo)

ggplot/geoms/geom_point.py

@@ -15,12 +15,13 @@ class geom_point(geom):
 
     def _plot_unit(self, pinfo, ax):
 
-        _abscent = {None: pinfo['color'], False: ''}
-        try:
-            if pinfo['facecolor'] in _abscent:
-                pinfo['facecolor'] = _abscent[pinfo['facecolor']]
-        except TypeError:
-            pass
+        fc = pinfo['facecolor']
+        if fc is None:
+            # default to color
+            pinfo['facecolor'] = pinfo['color']
+        elif fc is False:
+            # Matlab expects empty string instead of False
+            pinfo['facecolor'] = ''
 
         # for some reason, scatter doesn't default to the same color styles
         # as the axes.color_cycle

ggplot/geoms/geom_pointrange.py

@@ -0,0 +1,133 @@
+from __future__ import (absolute_import, division, print_function,
+                        unicode_literals)
+
+import sys
+
+import matplotlib as mpl
+from .geom import geom
+from ggplot.utils import is_categorical
+import numpy as np
+
+
+class geom_pointrange(geom):
+    """Plot intervals represented by vertical lines with a point in each interval
+
+     Parameters
+     ---------
+
+     x
+         x values of data
+     y
+         y value of the point for each x
+     ymin
+         lower end of the interval for each x
+     ymax
+         upper end of the interval for each x
+     alpha : float
+         alpha value, defaults to 1
+     color : string
+         line color, defaults to 'black'
+     fill : string
+         Fill type for the points, defaults to 'None'
+     linetype : string
+         line type, defaults to 'solid'
+     shape : string
+         shape of the points, defaults to 'o' (i.e. circles)
+     size : string
+         width of the line and size of the point, defaults to 2
+
+     Examples
+     --------
+
+     .. plot::
+         :include-source:
+
+         import numpy as np
+         import pandas as pd
+         from ggplot import *
+
+         np.random.seed(42)
+         x = np.linspace(0.5, 9.5, num=10)
+         y = np.random.randn(10)
+         ymin = y - np.random.uniform(0,1, size=10)
+         ymax = y + np.random.uniform(0,1, size=10)
+
+         data = pd.DataFrame({'x': x, 'y': y, 'ymin': ymin, 'ymax': ymax})
+
+         ggplot(aes(x='x', y='y', ymin='ymin', ymax='ymax'), data) \
+             + geom_pointrange()
+
+    """
+
+    DEFAULT_AES = {'alpha': 1, 'color': 'black', 'fill': None,
+                   'linetype': 'solid',
+                   'shape': 'o',
+                   'size': 2}
+    REQUIRED_AES = {'x', 'y', 'ymin', 'ymax'}
+    DEFAULT_PARAMS = {'stat': 'identity', 'position': 'identity', 'cmap': None}
+
+    _aes_renames = {'size': 'linewidth', 'linetype': 'linestyle', 'shape': 'marker', 'fill': 'facecolor'}
+    _units = {'alpha', 'color', 'linestyle', 'marker'}
+
+    def __init__(self, *args, **kwargs):
+        super(geom_pointrange, self).__init__(*args, **kwargs)
+        self._warning_printed = False
+
+    def _plot_unit(self, pinfo, ax):
+        # If x is categorical, calculate positions to plot
+        categorical = is_categorical(pinfo['x'])
+        if categorical:
+            x = pinfo.pop('x')
+            new_x = np.arange(len(x))
+            ax.set_xticks(new_x)
+            ax.set_xticklabels(x)
+            pinfo['x'] = new_x
+
+        if 'linewidth' in pinfo and isinstance(pinfo['linewidth'], list):
+            # ggplot also supports aes(size=...) but the current mathplotlib
+            # is not. See https://github.com/matplotlib/matplotlib/issues/2658
+            pinfo['linewidth'] = 4
+            if not self._warning_printed:
+                msg = "'geom_line()' currenty does not support the mapping of " +\
+                      "size ('aes(size=<var>'), using size=4 as a replacement.\n" +\
+                      "Use 'geom_line(size=x)' to set the size for the whole line.\n"
+                sys.stderr.write(msg)
+                self._warning_printed = True
+
+        # Plotting the line
+        pinfoline = dict(pinfo)
+        del pinfoline['marker']
+        del pinfoline['facecolor']
+        del pinfoline['y']
+
+        x = pinfoline.pop('x')
+        x = np.vstack([x, x])
+
+        ymin = pinfoline.pop('ymin')
+        ymax = pinfoline.pop('ymax')
+        y = np.vstack([ymin, ymax])
+
+        ax.plot(x, y, **pinfoline)
+
+        # Plotting the points
+        pinfopoint = dict(pinfo)
+        del pinfopoint['ymin']
+        del pinfopoint['ymax']
+        del pinfopoint['linestyle']
+
+        fc = pinfopoint['facecolor']
+        if fc is None:
+            # default to color
+            pinfopoint['facecolor'] = pinfopoint['color']
+        elif fc is False:
+            # Matlab expects empty string instead of False
+            pinfopoint['facecolor'] = ''
+
+        # for some reason, scatter doesn't default to the same color styles
+        # as the axes.color_cycle
+        if "color" not in pinfopoint and self.params['cmap'] is None:
+            pinfopoint["color"] = mpl.rcParams.get("axes.color_cycle", ["#333333"])[0]
+
+        pinfopoint['s'] = pinfopoint.pop('linewidth')**2*4
+
+        ax.scatter(**pinfopoint)

ggplot/tests/__init__.py

@@ -34,6 +34,8 @@ def teardown_package():
     'ggplot.tests.test_geom_bar',
     'ggplot.tests.test_qplot',
     'ggplot.tests.test_geom_lines',
+    'ggplot.tests.test_geom_linerange',
+    'ggplot.tests.test_geom_pointrange',
     'ggplot.tests.test_faceting',
     'ggplot.tests.test_stat_function',
     'ggplot.tests.test_scale_facet_wrap',
@@ -101,11 +103,14 @@ def make_test_fn(fname, purpose):
         base, ext = os.path.splitext(fname)
         return '%s-%s%s' % (base, purpose, ext)
     expected_fname = make_test_fn(actual_fname, 'expected')
+    # Save the figure before testing whether the original image
+    # actually exists. This make creating new tests much easier,
+    # as the result image can afterwards just be copied.
+    fig.savefig(actual_fname)
     if os.path.exists(orig_expected_fname):
         shutil.copyfile(orig_expected_fname, expected_fname)
     else:
         raise Exception("Baseline image %s is missing" % orig_expected_fname)
-    fig.savefig(actual_fname)
     err = compare_images(expected_fname, actual_fname,
                          tol, in_decorator=True)
     if err:

ggplot/tests/baseline_images/test_geom_linerange/geom_linerange.png

@@ -0,0 +1,29 @@
+from __future__ import (absolute_import, division, print_function,
+                        unicode_literals)
+
+from . import get_assert_same_ggplot, cleanup
+assert_same_ggplot = get_assert_same_ggplot(__file__)
+
+from ggplot import *
+
+import numpy as np
+import pandas as pd
+
+
+def _build_testing_df():
+    rst = np.random.RandomState(42)
+    x = np.linspace(0.5, 9.5, num=10)
+    y = rst.randn(10)
+    ymin = y - rst.uniform(0, 1, size=10)
+    ymax = y + rst.uniform(0, 1, size=10)
+
+    df = pd.DataFrame({'x': x, 'y': y, 'ymin': ymin, 'ymax': ymax})
+    return df
+
+
+@cleanup
+def test_geom_linerange():
+    df = _build_testing_df()
+    gg = ggplot(aes(x="x", y="y", ymin="ymin", ymax="ymax"), data=df)
+    gg = gg + geom_linerange()
+    assert_same_ggplot(gg, "geom_linerange")

ggplot/tests/baseline_images/test_geom_pointrange/geom_pointrange.png

@@ -0,0 +1,29 @@
+from __future__ import (absolute_import, division, print_function,
+                        unicode_literals)
+
+from . import get_assert_same_ggplot, cleanup
+assert_same_ggplot = get_assert_same_ggplot(__file__)
+
+from ggplot import *
+
+import numpy as np
+import pandas as pd
+
+
+def _build_testing_df():
+    rst = np.random.RandomState(42)
+    x = np.linspace(0.5, 9.5, num=10)
+    y = rst.randn(10)
+    ymin = y - rst.uniform(0, 1, size=10)
+    ymax = y + rst.uniform(0, 1, size=10)
+
+    df = pd.DataFrame({'x': x, 'y': y, 'ymin': ymin, 'ymax': ymax})
+    return df
+
+
+@cleanup
+def test_geom_pointrange():
+    df = _build_testing_df()
+    gg = ggplot(aes(x="x", y="y", ymin="ymin", ymax="ymax"), data=df)
+    gg = gg + geom_pointrange()
+    assert_same_ggplot(gg, "geom_pointrange")