colors: Add power-law normalization

bgamari · bgamari · commit 1b9fc6c2b556 · 2014-03-10T12:26:57.000-04:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -18,6 +18,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/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -52,6 +52,7 @@
 import six
 from six.moves import map, zip
 
+import warnings
 import re
 import numpy as np
 from numpy import ma
@@ -1148,6 +1149,80 @@ 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.
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -55,6 +55,19 @@ def test_LogNorm():
     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)
