added 'lightsource' class to colors module, shading_example.py

to illustrate it's usage to create shaded relief maps with imshow.

svn path=/trunk/matplotlib/; revision=6975

Author: Jeff Whitaker

CHANGELOG

@@ -1,3 +1,7 @@
+2009-03-14 Added 'lightsource' class to colors module for 
+           creating shaded relief maps.  shading_example.py
+           added to illustrate usage. - JSW
+
 2009-03-11 Ensure wx version >= 2.8; thanks to Sandro Tosi and
            Chris Barker. - EF
 

examples/pylab_examples/shading_example.py

@@ -0,0 +1,28 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.colors import lightsource
+
+# example showing how to make shaded relief plots 
+# like mathematica
+# (http://reference.wolfram.com/mathematica/ref/ReliefPlot.html )
+# or Generic Mapping Tools
+# (http://gmt.soest.hawaii.edu/gmt/doc/gmt/html/GMT_Docs/node145.html)
+
+# test data
+X,Y=np.mgrid[-5:5:0.1,-5:5:0.1]
+Z=X+np.sin(X**2+Y**2)
+# creat light source object.
+ls = lightsource(azdeg=270,altdeg=20)
+# shade data, creating an rgb array.
+rgb = ls.shade(Z,plt.cm.copper)
+# plot un-shaded and shaded images.
+plt.figure(figsize=(12,5))
+plt.subplot(121)
+plt.imshow(Z,cmap=plt.cm.copper)
+plt.title('imshow')
+plt.xticks([]); plt.yticks([])
+plt.subplot(122)
+plt.imshow(rgb)
+plt.title('imshow with shading')
+plt.xticks([]); plt.yticks([])
+plt.show()

lib/matplotlib/colors.py

@@ -900,3 +900,130 @@ def inverse(self, value):
 # compatibility with earlier class names that violated convention:
 normalize = Normalize
 no_norm = NoNorm
+
+def rgb_to_hsv(arr):
+    """
+    convert rgb values in a numpy array to hsv values
+    input and output arrays should have shape (M,N,3)
+    """
+    out = np.empty_like(arr)
+    arr_max = arr.max(-1)
+    delta = arr.ptp(-1)
+    s = delta / arr_max
+    s[delta==0] = 0
+    # red is max
+    idx = (arr[:,:,0] == arr_max)
+    out[idx, 0] = (arr[idx, 1] - arr[idx, 2]) / delta[idx]
+    # green is max
+    idx = (arr[:,:,1] == arr_max)
+    out[idx, 0] = 2. + (arr[idx, 2] - arr[idx, 0] ) / delta[idx]
+    # blue is max
+    idx = (arr[:,:,2] == arr_max)
+    out[idx, 0] = 4. + (arr[idx, 0] - arr[idx, 1] ) / delta[idx]
+    out[:,:,0] = (out[:,:,0]/6.0) % 1.0
+    out[:,:,1] = s
+    out[:,:,2] = arr_max
+    return out
+
+def hsv_to_rgb(hsv):
+    """
+    convert hsv values in a numpy array to rgb values
+    both input and output arrays have shape (M,N,3)
+    """
+    h = hsv[:,:,0]; s = hsv[:,:,1]; v = hsv[:,:,2]
+    r = np.empty_like(h); g = np.empty_like(h); b = np.empty_like(h)
+    i = (h*6.0).astype(np.int)
+    f = (h*6.0) - i
+    p = v*(1.0 - s)
+    q = v*(1.0 - s*f)
+    t = v*(1.0 - s*(1.0-f))
+    idx = i%6 == 0
+    r[idx] = v[idx]; g[idx] = t[idx]; b[idx] = p[idx]
+    idx = i == 1
+    r[idx] = q[idx]; g[idx] = v[idx]; b[idx] = p[idx]
+    idx = i == 2
+    r[idx] = p[idx]; g[idx] = v[idx]; b[idx] = t[idx]
+    idx = i == 3
+    r[idx] = p[idx]; g[idx] = q[idx]; b[idx] = v[idx]
+    idx = i == 4
+    r[idx] = t[idx]; g[idx] = p[idx]; b[idx] = v[idx]
+    idx = i == 5
+    r[idx] = v[idx]; g[idx] = p[idx]; b[idx] = q[idx]
+    idx = s == 0
+    r[idx] = v[idx]; g[idx] = v[idx]; b[idx] = v[idx]
+    return np.array((r,g,b)).T
+
+class lightsource(object):
+    """
+    Create a light source coming from the specified azimuth and elevation.
+    Angles are in degrees, with the azimuth measured
+    clockwise from south and elevation up from the zero plane of the surface.
+    The :meth:`shade` is used to produce rgb values for a shaded relief image
+    given a data array.
+    """
+    def __init__(self,azdeg=315,altdeg=45,\
+                 hsv_min_val=0,hsv_max_val=1,hsv_min_sat=1,hsv_max_sat=0):
+       """
+       Specify the azimuth (measured clockwise from south) and altitude 
+       (measured up from the plane of the surface) of the light source
+       in degrees.
+
+       The color of the resulting image will be darkened
+       by moving the (s,v) values (in hsv colorspace) toward
+       (hsv_min_sat, hsv_min_val) in the shaded regions, or
+       lightened by sliding (s,v) toward
+       (hsv_max_sat hsv_max_val) in regions that are illuminated.
+       The default extremes are chose so that completely shaded points
+       are nearly black (s = 1, v = 0) and completely illuminated points
+       are nearly white (s = 0, v = 1). 
+       """
+       self.azdeg = azdeg
+       self.altdeg = altdeg
+       self.hsv_min_val = hsv_min_val
+       self.hsv_max_val = hsv_max_val
+       self.hsv_min_sat = hsv_min_sat
+       self.hsv_max_sat = hsv_max_sat
+
+    def shade(self,data,cmap):
+        """
+        Take the input data array, convert to HSV values in the
+        given colormap, then adjust those color values 
+        to given the impression of a shaded relief map with a
+        specified light source.
+        RGBA values are returned, which can then be used to 
+        plot the shaded image with imshow.
+        """
+        # imagine an artificial sun placed at infinity in
+        # some azimuth and elevation position illuminating our surface. The parts of
+        # the surface that slope toward the sun should brighten while those sides
+        # facing away should become darker.
+        # convert alt, az to radians
+        az = self.azdeg*np.pi/180.0
+        alt = self.altdeg*np.pi/180.0
+        # gradient in x and y directions
+        dx, dy = np.gradient(data)
+        slope = 0.5*np.pi - np.arctan(np.hypot(dx, dy))
+        aspect = np.arctan2(dx, dy)
+        intensity = np.sin(alt)*np.sin(slope) + np.cos(alt)*np.cos(slope)*np.cos(-az -\
+                aspect - 0.5*np.pi)
+        # rescale to interval -1,1
+        # +1 means maximum sun exposure and -1 means complete shade.
+        intensity = (intensity - intensity.min())/(intensity.max() - intensity.min())
+        intensity = 2.*intensity - 1.
+        # convert to rgb, then rgb to hsv
+        rgb = cmap((data-data.min())/(data.max()-data.min()))
+        hsv = rgb_to_hsv(rgb[:,:,0:3])
+        # modify hsv values to simulate illumination.
+        hsv[:,:,1] = np.where(np.logical_and(np.abs(hsv[:,:,1])>1.e-10,intensity>0),\
+                (1.-intensity)*hsv[:,:,1]+intensity*self.hsv_max_sat, hsv[:,:,1])
+        hsv[:,:,2] = np.where(intensity > 0, (1.-intensity)*hsv[:,:,1] +\
+                intensity*self.hsv_max_val, hsv[:,:,2])
+        hsv[:,:,1] = np.where(np.logical_and(np.abs(hsv[:,:,1])>1.e-10,intensity<0),\
+                (1.+intensity)*hsv[:,:,1]-intensity*self.hsv_min_sat, hsv[:,:,1])
+        hsv[:,:,2] = np.where(intensity < 0, (1.+intensity)*hsv[:,:,1] -\
+                intensity*self.hsv_min_val, hsv[:,:,2])
+        hsv[:,:,1:] = np.where(hsv[:,:,1:]<0.,0,hsv[:,:,1:])
+        hsv[:,:,1:] = np.where(hsv[:,:,1:]>1.,1,hsv[:,:,1:])
+        # convert modified hsv back to rgb.
+        rgb[:,:,0:3] = hsv_to_rgb(hsv)
+        return rgb