Merge pull request #1204 from bgamari/power-law-norm

tacaswell · tacaswell · commit d0250b4dbc34 · 2014-04-03T22:27:52.000-04:00
Add power-law normalization
diff --git a/CHANGELOG b/CHANGELOG
@@ -31,6 +31,9 @@
            all pyplot.tri* methods) and mlab.griddata.  Deprecated
            matplotlib.delaunay module. - IMT
 
+2013-11-05 Add power-law normalization method. This is useful for,
+	   e.g., showing small populations in a "hist2d" histogram.
+
 2013-10-27 Added get_rlabel_position and set_rlabel_position methods to
            PolarAxes to control angular position of radial tick labels.
 
diff --git a/doc/users/whats_new.rst b/doc/users/whats_new.rst
@@ -32,6 +32,14 @@ Phil Elson rewrote of the documentation and userguide for both Legend and PathEf
 New plotting features
 ---------------------
 
+Power-law normalization
+```````````````````````
+Ben Gamari added a power-law normalization method,
+:class:`~matplotlib.colors.PowerNorm`. This class maps a range of 
+values to the interval [0,1] with power-law scaling with the exponent
+provided by the constructor's `gamma` argument. Power law normalization
+can be useful for, e.g., emphasizing small populations in a histogram.
+
 Fully customizable boxplots
 ````````````````````````````
 Paul Hobson overhauled the :func:`~matplotlib.pyplot.boxplot` method such
diff --git a/examples/api/power_norm_demo.py b/examples/api/power_norm_demo.py
@@ -0,0 +1,27 @@
+#!/usr/bin/python
+
+from matplotlib import pyplot as plt
+import matplotlib.colors as mcolors
+import numpy as np
+from numpy.random import multivariate_normal
+
+data = np.vstack([multivariate_normal([10, 10], [[3, 5],[4, 2]], size=100000),
+                  multivariate_normal([30, 20], [[2, 3],[1, 3]], size=1000)
+                 ])
+
+gammas = [0.8, 0.5, 0.3]
+xgrid = np.floor((len(gammas) + 1.) / 2)
+ygrid = np.ceil((len(gammas) + 1.) / 2)
+
+plt.subplot(xgrid, ygrid, 1)
+plt.title('Linear normalization')
+plt.hist2d(data[:,0], data[:,1], bins=100)
+
+for i, gamma in enumerate(gammas):
+    plt.subplot(xgrid, ygrid, i + 2)
+    plt.title('Power law normalization\n$(\gamma=%1.1f)$' % gamma)
+    plt.hist2d(data[:, 0], data[:, 1],
+               bins=100, norm=mcolors.PowerNorm(gamma))
+
+plt.subplots_adjust(hspace=0.39)
+plt.show()
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5839,7 +5839,9 @@ def hist2d(self, x, y, bins=10, range=None, normed=False, weights=None,
         -----
         Rendering the histogram with a logarithmic color scale is
         accomplished by passing a :class:`colors.LogNorm` instance to
-        the *norm* keyword argument.
+        the *norm* keyword argument. Likewise, power-law normalization
+        (similar in effect to gamma correction) can be accomplished with
+        :class:`colors.PowerNorm`.
 
         Examples
         --------
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -52,6 +52,7 @@
 import six
 from six.moves import zip
 
+import warnings
 import re
 import numpy as np
 from numpy import ma
@@ -617,24 +618,24 @@ def __call__(self, X, alpha=None, bytes=False):
         return rgba
 
     def set_bad(self, color='k', alpha=None):
-        '''Set color to be used for masked values.
-        '''
+        """Set color to be used for masked values.
+        """
         self._rgba_bad = colorConverter.to_rgba(color, alpha)
         if self._isinit:
             self._set_extremes()
 
     def set_under(self, color='k', alpha=None):
-        '''Set color to be used for low out-of-range values.
+        """Set color to be used for low out-of-range values.
            Requires norm.clip = False
-        '''
+        """
         self._rgba_under = colorConverter.to_rgba(color, alpha)
         if self._isinit:
             self._set_extremes()
 
     def set_over(self, color='k', alpha=None):
-        '''Set color to be used for high out-of-range values.
+        """Set color to be used for high out-of-range values.
            Requires norm.clip = False
-        '''
+        """
         self._rgba_over = colorConverter.to_rgba(color, alpha)
         if self._isinit:
             self._set_extremes()
@@ -651,7 +652,7 @@ def _set_extremes(self):
         self._lut[self._i_bad] = self._rgba_bad
 
     def _init(self):
-        '''Generate the lookup table, self._lut'''
+        """Generate the lookup table, self._lut"""
         raise NotImplementedError("Abstract class only")
 
     def is_gray(self):
@@ -937,9 +938,9 @@ def inverse(self, value):
             return vmin + value * (vmax - vmin)
 
     def autoscale(self, A):
-        '''
+        """
         Set *vmin*, *vmax* to min, max of *A*.
-        '''
+        """
         self.vmin = ma.min(A)
         self.vmax = ma.max(A)
 
@@ -1008,9 +1009,9 @@ def inverse(self, value):
             return vmin * pow((vmax / vmin), value)
 
     def autoscale(self, A):
-        '''
+        """
         Set *vmin*, *vmax* to min, max of *A*.
-        '''
+        """
         A = ma.masked_less_equal(A, 0, copy=False)
         self.vmin = ma.min(A)
         self.vmax = ma.max(A)
@@ -1140,8 +1141,82 @@ def autoscale_None(self, A):
         self._transform_vmin_vmax()
 
 
+class PowerNorm(Normalize):
+    """
+    Normalize a given value to the ``[0, 1]`` interval with a power-law
+    scaling. This will clip any negative data points to 0.
+    """
+    def __init__(self, gamma, vmin=None, vmax=None, clip=False):
+        Normalize.__init__(self, vmin, vmax, clip)
+        self.gamma = gamma
+
+    def __call__(self, value, clip=None):
+        if clip is None:
+            clip = self.clip
+
+        result, is_scalar = self.process_value(value)
+
+        self.autoscale_None(result)
+        gamma = self.gamma
+        vmin, vmax = self.vmin, self.vmax
+        if vmin > vmax:
+            raise ValueError("minvalue must be less than or equal to maxvalue")
+        elif vmin == vmax:
+            result.fill(0)
+        else:
+            if clip:
+                mask = ma.getmask(result)
+                val = ma.array(np.clip(result.filled(vmax), vmin, vmax),
+                                mask=mask)
+            resdat = result.data
+            resdat -= vmin
+            np.power(resdat, gamma, resdat)
+            resdat /= (vmax - vmin) ** gamma
+            result = np.ma.array(resdat, mask=result.mask, copy=False)
+            result[value < 0] = 0
+        if is_scalar:
+            result = result[0]
+        return result
+
+    def inverse(self, value):
+        if not self.scaled():
+            raise ValueError("Not invertible until scaled")
+        gamma = self.gamma
+        vmin, vmax = self.vmin, self.vmax
+
+        if cbook.iterable(value):
+            val = ma.asarray(value)
+            return ma.power(value, 1. / gamma) * (vmax - vmin) + vmin
+        else:
+            return pow(value, 1. / gamma) * (vmax - vmin) + vmin
+
+    def autoscale(self, A):
+        """
+        Set *vmin*, *vmax* to min, max of *A*.
+        """
+        self.vmin = ma.min(A)
+        if self.vmin < 0:
+            self.vmin = 0
+            warnings.warn("Power-law scaling on negative values is "
+                          "ill-defined, clamping to 0.")
+
+        self.vmax = ma.max(A)
+
+    def autoscale_None(self, A):
+        ' autoscale only None-valued vmin or vmax'
+        if self.vmin is None and np.size(A) > 0:
+            self.vmin = ma.min(A)
+            if self.vmin < 0:
+                self.vmin = 0
+                warnings.warn("Power-law scaling on negative values is "
+                              "ill-defined, clamping to 0.")
+
+        if self.vmax is None and np.size(A) > 0:
+            self.vmax = ma.max(A)
+
+
 class BoundaryNorm(Normalize):
-    '''
+    """
     Generate a colormap index based on discrete intervals.
 
     Unlike :class:`Normalize` or :class:`LogNorm`,
@@ -1152,9 +1227,9 @@ class BoundaryNorm(Normalize):
     piece-wise linear interpolation, but using integers seems
     simpler, and reduces the number of conversions back and forth
     between integer and floating point.
-    '''
+    """
     def __init__(self, boundaries, ncolors, clip=False):
-        '''
+        """
         *boundaries*
             a monotonically increasing sequence
         *ncolors*
@@ -1171,7 +1246,7 @@ def __init__(self, boundaries, ncolors, clip=False):
         Out-of-range values are mapped to -1 if low and ncolors
         if high; these are converted to valid indices by
         :meth:`Colormap.__call__` .
-        '''
+        """
         self.clip = clip
         self.vmin = boundaries[0]
         self.vmax = boundaries[-1]
@@ -1209,11 +1284,11 @@ def inverse(self, value):
 
 
 class NoNorm(Normalize):
-    '''
+    """
     Dummy replacement for Normalize, for the case where we
     want to use indices directly in a
     :class:`~matplotlib.cm.ScalarMappable` .
-    '''
+    """
     def __call__(self, value, clip=None):
         return value
 
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -47,14 +47,27 @@ def test_BoundaryNorm():
 
 def test_LogNorm():
     """
-    LogNorm igornoed clip, now it has the same
+    LogNorm ignored clip, now it has the same
     behavior as Normalize, e.g., values > vmax are bigger than 1
     without clip, with clip they are 1.
     """
     ln = mcolors.LogNorm(clip=True, vmax=5)
     assert_array_equal(ln([1, 6]), [0, 1.0])
 
 
+def test_PowerNorm():
+    a = np.array([0, 0.5, 1, 1.5], dtype=np.float)
+    pnorm = mcolors.PowerNorm(1)
+    norm = mcolors.Normalize()
+    assert_array_almost_equal(norm(a), pnorm(a))
+
+    a = np.array([-0.5, 0, 2, 4, 8], dtype=np.float)
+    expected = [0, 0, 1./16, 1./4, 1]
+    pnorm = mcolors.PowerNorm(2, vmin=0, vmax=8)
+    assert_array_almost_equal(pnorm(a), expected)
+    assert_array_almost_equal(a[1:], pnorm.inverse(pnorm(a))[1:])
+
+
 def test_Normalize():
     norm = mcolors.Normalize()
     vals = np.arange(-10, 10, 1, dtype=np.float)
