examples and better spines

Author: abarbu

src/Graphics/Matplotlib.hs

@@ -52,16 +52,19 @@
 -- Right now there's no easy way to bind to an option other than the last one
 -- unless you want to pass options in as parameters.
 --
--- TODO The generated Python code should follow some invariants. It must maintain the
+-- The generated Python code should follow some invariants. It must maintain the
 -- current figure in "fig", all available axes in "axes", and the current axis
--- in "ax".
+-- in "ax". Plotting commands should use the current axis, never the plot
+-- itself; the two APIs are almost identical. When creating low-level bindings
+-- one must remember to call "plot.sci" to set the current image when plotting a
+-- graph. The current spine of the axes that's being manipulated is in "spine".
 -----------------------------------------------------------------------------
 
 module Graphics.Matplotlib
   ( module Graphics.Matplotlib
     -- * Creating custom plots and applying options
-  , Matplotlib(), Option(),(@@), (%), o1, o2, (##), (#), mp, def, readData,
-    str, raw, lit, updateAxes, updateFigure)
+  , Matplotlib(), Option(),(@@), (%), o1, o2, (##), (#), mp, def, readData
+  , str, raw, lit, updateAxes, updateFigure, mapLinear)
 where
 import Data.List
 import Data.Aeson
@@ -101,20 +104,21 @@ histogram values bins = readData [values] % dataHistogram 0 bins
 
 -- | Plot a 2D histogram for the given values with 'bins'
 histogram2D x y = readData [x,y] %
-  mp # "plot.hist2d(data[0], data[1]" ## ")"
+  mp # "plot.sci(ax.hist2d(data[0], data[1]" ## ")[-1])"
 
 -- | Plot the given values as a scatter plot
 scatter :: (ToJSON t1, ToJSON t) => t1 -> t -> Matplotlib
 scatter x y = readData (x, y)
-  % mp # "ax.scatter(data[0], data[1]" ## ")"
+  % mp # "plot.sci(ax.scatter(data[0], data[1]" ## "))"
 
 -- | Plot a line
 line :: (ToJSON t1, ToJSON t) => t1 -> t -> Matplotlib
 line x y = plot x y `def` [o1 "-"]
 
 -- | Like 'plot' but takes an error bar value per point
-errorbar xs ys errs = readData (xs, ys, errs)
-  % mp # "ax.errorbar(data[0], data[1], yerr=data[2]" ## ")"
+-- errorbar :: (ToJSON x, ToJSON y, ToJSON xs, ToJSON ys) => x -> y -> Maybe xs -> Maybe ys -> Matplotlib
+errorbar xs ys xerrs yerrs = readData (xs, ys, xerrs, yerrs)
+  % mp # "ax.errorbar(data[0], data[1], xerr=data[2], yerr=data[3]" ## ")"
 
 -- | Plot a line given a function that will be executed for each element of
 -- given list. The list provides the x values, the function the y values.
@@ -165,12 +169,12 @@ line1 y = line [0..length y] y
 -- | Plot a matrix
 matShow :: ToJSON a => a -> Matplotlib
 matShow d = readData d
-            % (mp # "plot.matshow(data" ## ")")
+            % (mp # "plot.sci(ax.matshow(data" ## "))")
 
 -- | Plot a matrix
 pcolor :: ToJSON a => a -> Matplotlib
 pcolor d = readData d
-            % (mp # "plot.pcolor(np.array(data)" ## ")")
+            % (mp # "plot.sci(ax.pcolor(np.array(data)" ## "))")
 
 -- | Plot a KDE of the given functions; a good bandwith will be chosen automatically
 density :: [Double] -> Maybe (Double, Double) -> Matplotlib
@@ -197,7 +201,7 @@ setUnicode b = mp # "matplotlib.rcParams['text.latex.unicode'] = " # b
 
 -- | Plot the 'a' and 'b' entries of the data object
 dataPlot :: (MplotValue val, MplotValue val1) => val1 -> val -> Matplotlib
-dataPlot a b = mp # "p = plot.plot(data[" # a # "], data[" # b # "]" ## ")"
+dataPlot a b = mp # "p = ax.plot(data[" # a # "], data[" # b # "]" ## ")"
 
 -- | Plot the Haskell objects 'x' and 'y' as a line
 plot :: (ToJSON t, ToJSON t1) => t1 -> t -> Matplotlib
@@ -216,7 +220,7 @@ dateLine x y xunit (yearStart, monthStart, dayStart) =
 
 -- | Create a histogram for the 'a' entry of the data array
 dataHistogram :: (MplotValue val1, MplotValue val) => val1 -> val -> Matplotlib
-dataHistogram a bins = mp # "plot.hist(data[" # a # "]," # bins ## ")"
+dataHistogram a bins = mp # "ax.hist(data[" # a # "]," # bins ## ")"
 
 -- | Create a scatter plot accessing the given fields of the data array
 dataScatter :: (MplotValue val1, MplotValue val) => val1 -> val -> Matplotlib
@@ -309,8 +313,14 @@ acorr x = readData x % mp # "ax.acorr(data" ## ")"
 -- | Plot text at a specified location
 text x y s = mp # "ax.text(" # x # "," # y # "," # raw s ## ")"
 
+figText x y s = mp # "plot.figtext(" # x # "," # y # "," # raw s ## ")"
+
 -- * Layout, axes, and legends
 
+-- | Square up the aspect ratio of a plot.
+setAspect :: Matplotlib
+setAspect = mp # "ax.set_aspect(" ## ")"
+
 -- | Square up the aspect ratio of a plot.
 squareAxes :: Matplotlib
 squareAxes = mp # "ax.set_aspect('equal')"
@@ -353,14 +363,6 @@ legend = mp # "ax.legend(" ## ")"
 -- TODO This refers to the plot and not an axis. Might cause trouble with subplots
 colorbar = mp # "plot.colorbar(" ## ")"
 
--- | Set the spacing of ticks on the x axis
-axisXTickSpacing :: (MplotValue val1, MplotValue val) => val1 -> val -> Matplotlib
-axisXTickSpacing nr width = mp # "ax.set_xticks(np.arange(" # nr # ")+" # width ## ")"
-
--- | Set the labels on the x axis
-axisXTickLabels :: MplotValue val => val -> Matplotlib
-axisXTickLabels labels = mp # "ax.set_xticklabels( (" # labels # ") " ## " )"
-
 -- | Add a title
 title :: String -> Matplotlib
 title s = mp # "ax.set_title(" # raw s ## ")"
@@ -397,22 +399,70 @@ zLabel label = mp # "ax.set_zlabel(" # raw label ## ")"
 
 setSizeInches w h = mp # "fig.set_size_inches(" # w # "," # h # ", forward=True)"
 
+tightLayout = mp # "fig.tight_layout()"
+
+xkcd = mp # "plot.xkcd()"
+
+-- * Ticks
+
+xticks l = mp # "ax.set_xticks(" # l # ")"
+yticks l = mp # "ax.set_yticks(" # l # ")"
+zticks l = mp # "ax.set_zticks(" # l # ")"
+
+xtickLabels l = mp # "ax.set_xticklabels(" # l # ")"
+ytickLabels l = mp # "ax.set_yticklabels(" # l # ")"
+ztickLabels l = mp # "ax.set_zticklabels(" # l # ")"
+
+-- | Set the spacing of ticks on the x axis
+axisXTickSpacing :: (MplotValue val1, MplotValue val) => val1 -> val -> Matplotlib
+axisXTickSpacing nr width = mp # "ax.set_xticks(np.arange(" # nr # ")+" # width ## ")"
+
+-- | Set the labels on the x axis
+axisXTickLabels :: MplotValue val => val -> Matplotlib
+axisXTickLabels labels = mp # "ax.set_xticklabels( (" # labels # ") " ## " )"
+
+-- | Set the spacing of ticks on the y axis
+axisYTickSpacing :: (MplotValue val1, MplotValue val) => val1 -> val -> Matplotlib
+axisYTickSpacing nr width = mp # "ax.set_yticks(np.arange(" # nr # ")+" # width ## ")"
+
+-- | Set the labels on the y axis
+axisYTickLabels :: MplotValue val => val -> Matplotlib
+axisYTickLabels labels = mp # "ax.set_yticklabels( (" # labels # ") " ## " )"
+
+axisXTicksPosition p = mp # "ax.xaxis.set_ticks_position('" # p # "')"
+axisYTicksPosition p = mp # "ax.yaxis.set_ticks_position('" # p # "')"
+
+-- * Spines
+
+spine s = mp # "spine = ax.spines['" # s # "']"
+
+spineSetBounds l h = mp # "spine.set_bounds(" # l # "," # h # ")"
+
+spineSetVisible b = mp # "spine.set_visible(" # b # ")"
+
+spineSetPosition s n = mp # "spine.set_position((" # s # "," # n # "))"
+
 -- * Subplots
 
+setAx = mp # "plot.sca(ax) "
+
 -- | Create a subplot with the coordinates (r,c,f)
-addSubplot r c f = mp # "ax = plot.gcf().add_subplot(" # r # c # f ## ")" % updateAxes
+addSubplot r c f = mp # "ax = plot.gcf().add_subplot(" # r # c # f ## ")" % updateAxes % setAx
 
 -- | Access a subplot with the coordinates (r,c,f)
-getSubplot r c f = mp # "ax = plot.subplot(" # r # "," # c # "," # f ## ")" % updateAxes
+getSubplot r c f = mp # "ax = plot.subplot(" # r # "," # c # "," # f ## ")" % updateAxes % setAx
 
 -- | Creates subplots and stores them in an internal variable
 subplots = mp # "fig, axes = plot.subplots(" ## ")"
+  % mp # "axes = np.asarray(axes)"
+  % mp # "axes = axes.flatten()"
+  % updateAxes % setAx
 
 -- | Access a subplot
-setSubplot s = mp # "ax = axes[" # s # "]"
+setSubplot s = mp # "ax = axes[" # s # "]" % setAx
 
 -- | Add axes to a figure
-axes = mp # "ax = plot.axes(" ## ")" % updateAxes
+axes = mp # "ax = plot.axes(" ## ")" % updateAxes % setAx
 
 -- | Creates a new figure with the given id. If the Id is already in use it
 -- switches to that figure.

src/Graphics/Matplotlib/Internal.hs

@@ -164,13 +164,29 @@ instance MplotValue [Int] where
   toPython s = "[" ++ f s ++ "]"
     where f [] = ""
           f (x:xs) = toPython x ++ "," ++ f xs
+instance MplotValue [R] where
+  toPython s = "[" ++ f s ++ "]"
+    where f [] = ""
+          f (x:xs) = toPython x ++ "," ++ f xs
+instance MplotValue [S] where
+  toPython s = "[" ++ f s ++ "]"
+    where f [] = ""
+          f (x:xs) = toPython x ++ "," ++ f xs
+instance MplotValue [L] where
+  toPython s = "[" ++ f s ++ "]"
+    where f [] = ""
+          f (x:xs) = toPython x ++ "," ++ f xs
 instance MplotValue Bool where
   toPython s = show s
 instance (MplotValue x) => MplotValue (x, x) where
-  toPython (n, v) = toPython n ++ " = " ++ toPython v
+  toPython (k, v) = "(" ++ toPython k ++ ", " ++ toPython v ++ ")"
 instance (MplotValue (x, y)) => MplotValue [(x, y)] where
-  toPython [] = ""
-  toPython (x:xs) = toPython x ++ ", " ++ toPython xs
+  toPython s = "[" ++ f s ++ "]"
+    where f [] = ""
+          f (x:xs) = toPython x ++ "," ++ f xs
+instance MplotValue x => MplotValue (Maybe x) where
+  toPython Nothing  = "None"
+  toPython (Just x) = toPython x
 
 default (Integer, Int, Double)
 

test/Spec.hs

@@ -1,4 +1,4 @@
-{-# language ExtendedDefaultRules, ScopedTypeVariables, QuasiQuotes #-}
+{-# language ExtendedDefaultRules, ScopedTypeVariables, QuasiQuotes, ParallelListComp #-}
 
 import Test.Tasty
 import Test.Tasty.Runners
@@ -133,7 +133,7 @@ testPlotGolden name fn =
                                                        case (stderr, reads stderr) of
                                                          ("inf", _) -> return Nothing
                                                          (_, [(x :: Double, _)]) ->
-                                                           if x < 30 then
+                                                           if x < 35 then
                                                              return $ Just $ "Images very different; PSNR too low " ++ show x else
                                                              return Nothing)))
                                    (BS.writeFile ref))
@@ -171,6 +171,11 @@ basicTests f = testGroup "Basic tests"
   , f "eventplot" meventplot
   , f "errorbar" merrorbar
   , f "scatterhist" mscatterHist
+  , f "histMulti" mhistMulti
+  , f "spines" mspines
+  , f "hists" mhists
+  , f "hinton" mhinton
+  , f "integral" mintegral
   ]
 
 fragileTests f = testGroup "Fragile tests"
@@ -252,6 +257,7 @@ mlegend = plotMapLinear (\x -> x ** 2) 0 1 100 @@ [o2 "label" "x^2"]
 
 -- | http://matplotlib.org/examples/pylab_examples/hist2d_log_demo.html
 mhist2DLog = histogram2D x y @@ [o2 "bins" 40, o2 "norm" $ lit "mcolors.LogNorm()"]
+  % setAx
   % colorbar
   where (x:y:_) = chunksOf 10000 normals
 
@@ -262,10 +268,10 @@ meventplot = plot xs ys
   where xs = sort $ take 10 uniforms
         ys = map (\x -> x ** 2) xs
 
-merrorbar = errorbar xs ys errs @@ [o2 "errorevery" 2]
+merrorbar = errorbar xs ys (Nothing :: Maybe [Double]) (Just errs) @@ [o2 "errorevery" 2]
   where xs = [0.1,0.2..4]
         ys = map (\x -> exp $ -x) xs
-        errs = map (\x -> 0.1 + 0.1 * sqrt x) xs
+        errs = [map (\x -> 0.1 + 0.1 * sqrt x) xs, map (\x -> 0.1 + 0.1 * sqrt x) ys]
 
 mboxplot = subplots @@ [o2 "ncols" 2, o2 "sharey" True]
   % setSubplot "0"
@@ -280,8 +286,7 @@ mviolinplot = subplots @@ [o2 "ncols" 2, o2 "sharey" True]
   % violinplot (take 3 $ chunksOf 100 $ map (* 2) $ normals) @@ [o2 "showmeans" True, o2 "showmedians" True, o2 "vert" False]
 
 -- | http://matplotlib.org/examples/pylab_examples/scatter_hist.html
-mscatterHist = readData ()
-  % figure 0
+mscatterHist = figure 0
   % setSizeInches 8 8
   -- The scatter plot
   % axes @@ [o1 ([left, bottom', width, height] :: [Double])]
@@ -309,3 +314,92 @@ mscatterHist = readData ()
         xymax = maximum [maximum $ map abs x, maximum $ map abs y]
         lim = ((fromIntegral $ round $ xymax / binwidth) + 1) * binwidth
         bins = [-lim,-lim+binwidth..(lim + binwidth)]
+
+mhistMulti = subplots @@ [o2 "nrows" 2, o2 "ncols" 2]
+  % setSubplot 0
+  % histogram x nrBins @@ [o2 "normed" 1, o2 "histtype" "bar", o2 "color" ["red", "tan", "lime"], o2 "label" ["red", "tan", "lime"]]
+  % legend @@ [o2 "prop" $ lit "{'size': 10}"]
+  % title "bars with legend"
+  % setSubplot 1
+  % histogram x nrBins @@ [o2 "normed" 1, o2 "histtype" "bar", o2 "stacked" True]
+  % title "stacked bar"
+  % setSubplot 2
+  % histogram x nrBins @@ [o2 "histtype" "step", o2 "stacked" True, o2 "fill" False]
+  % title "stacked bar"
+  % setSubplot 3
+  % histogram (map (\x -> take x normals) [2000, 5000, 10000]) nrBins @@ [o2 "histtype" "bar"]
+  % title "different sample sizes"
+  % tightLayout
+  where nrBins = 10
+        x = take 3 $ chunksOf 1000 $ normals
+
+mspines = plot x y @@ [o1 "k--"]
+  % plot x y' @@ [o1 "ro"]
+  % xlim 0 (2 * pi)
+  % xticks [0 :: Double, pi, 2*pi]
+  % xtickLabels (map raw ["0", "$\\pi$", "2$\\pi$"])
+  % ylim (-1.5) 1.5
+  % yticks [-1 :: Double, 0, 1]
+  % spine "left"
+  % spineSetBounds (-1) 1
+  % spine "right"
+  % spineSetVisible False
+  % spine "top"
+  % spineSetVisible False
+  % axisYTicksPosition "left"
+  % axisXTicksPosition "bottom"
+  where x  = mapLinear (\x -> x) 0 (2 * pi) 50
+        y  = map sin x
+        y' = zipWith (\a b -> a + 0.1*b) y normals
+
+mhists = h 10 1.5
+       % h 4 1
+       % h 15 2
+       % h 6 0.5
+  where ns mu var = map (\x -> mu + x * var) $ take 1000 normals
+        h mu var = histogram (ns mu var) 25 @@ [o2 "histtype" "stepfilled"
+                                             ,o2 "alpha" 0.8
+                                             ,o2 "normed" True]
+
+mhinton = mp # "ax.patch.set_facecolor('gray')"
+  % setAspect @@ [o1 "equal", o1 "box"]
+  % mp # "ax.xaxis.set_major_locator(plot.NullLocator())"
+  % mp # "ax.yaxis.set_major_locator(plot.NullLocator())"
+  % foldl (\a (x,y,w) -> a % f x y w) mp m
+  % mp # "ax.autoscale_view()"
+  % mp # "ax.invert_yaxis()"
+  where m = [ (x,y,w) | x <- [0..19], y <- [0..19] | w <- (map (\x -> x - 0.5) normals) ]
+        maxWeight = maximum $ map (\(_,_,v) -> abs v) m
+        f x y w = mp # "ax.add_patch(plot.Rectangle("
+                     # "[" # (x - size / 2) # "," # (y - size / 2) # "]"
+                     # ", " # size # ", " # size
+                     # ", facecolor='" # color # "', edgecolor='" # color # "'))"
+          where color = if w > 0 then "white" else "black"
+                size  = sqrt $ abs w / maxWeight
+
+mintegral = subplots
+  % plot x y @@ [o1 "r", o2 "linewidth" 2]
+  % ylim 0 (maximum y)
+  % mp # "ax.add_patch(plot.Polygon(" # ([(a, 0)] ++ zip ix iy ++ [(b,0)]) ## "))"
+    @@ [o2 "facecolor" "0.9", o2 "edgecolor" "0.5"]
+  % text (0.5 * (a + b)) 30 [r|$\int_a^b f(x)\mathrm{d}x$|]
+    @@ [o2 "horizontalalignment" "center", o2 "fontsize" 20]
+  % figText 0.9 0.05 "$x$"
+  % figText 0.1 0.9  "$y$"
+  % spine "right"
+  % spineSetVisible False
+  % spine "top"
+  % spineSetVisible False
+  % axisXTicksPosition "bottom"
+  % xticks (a, b)
+  % xtickLabels (raw "$a$", raw "$b$")
+  % yticks ([] :: [Double])
+  where func x = (x - 3) * (x - 5) * (x - 7) + 85
+        -- integral limits
+        a = 2 :: Double
+        b = 9 :: Double
+        (x :: [Double]) = mapLinear (\x -> x) 0 10 100
+        y = map func x
+        -- shaded region
+        (ix :: [Double]) = mapLinear (\x -> x) a b 100
+        iy = map func ix