work in progress simple matlab-like plotting API with multiple backends

HackPlotLib.hs

{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

-- | Some plotting utilities.
module HackPlotLib
       ( Plotter
         -- * plotting functions
       , figure, figure'
       , plot, plot', plot3, plot3'
       , subplot, hold
       , xlabel, ylabel, zlabel
       , title, legend
         -- * run a plotter
       , runMatlab--, runMatplotlib
       ) where

-- import Control.Lens ( (^.) )
import Control.Monad.State.Lazy ( State )
import qualified Control.Monad.State.Lazy as State
import qualified Data.Foldable as F
import Data.List ( intercalate )
import Data.Sequence ( Seq, (|>) )
--import qualified Data.Sequence as Seq
import Data.Text.Lazy ( Text )
import Data.Text.Lazy.Builder ( Builder, fromString, singleton, toLazyText )
import Data.Text.Lazy.Builder.RealFloat ( realFloat )
import Text.Printf ( printf )

data PlotCommand =
    Figure (Maybe Int)
  | Plot [(Double, Double)] [String]
  | Plot3 [(Double, Double, Double)] [String]
  | Hold Bool
  | Subplot Int Int Int
  | XLabel String
  | YLabel String
  | ZLabel String
  | Title String
  | Legend [String]

data PlotterState = PlotterState (Seq PlotCommand)

-- | the plotter monad
newtype Plotter a = Plotter (State PlotterState a)
                 deriving ( Functor, Applicative, Monad
                          , State.MonadState PlotterState
                          )

command :: PlotCommand -> Plotter ()
command c = do
  PlotterState commands <- State.get
  State.put (PlotterState (commands |> c))

-- | new figure
figure :: Plotter ()
figure = command (Figure Nothing)

-- | new figure with specific number
figure' :: Int -> Plotter ()
figure' k = command (Figure (Just k))

-- | plot a list of (x,y) pairs
plot :: [(Double, Double)] -> Plotter ()
plot v = plot' v []

-- | plot a list of (x,y) pairs with options
plot' :: [(Double, Double)] -> [String] ->  Plotter ()
plot' v o = command (Plot v o)

-- | plot a list of (x,y,z) pairs
plot3 :: [(Double, Double, Double)] -> Plotter ()
plot3 v = plot3' v []

-- | plot a list of (x,y,z) pairs with options
plot3' :: [(Double, Double, Double)] -> [String] -> Plotter ()
plot3' v o = command (Plot3 v o)

-- | make a subplot
subplot :: Int -> Int -> Int -> Plotter ()
subplot kx ky kz = command (Subplot kx ky kz)

-- | change the hold state
hold :: Bool -> Plotter ()
hold = command . Hold

-- | X axis label
xlabel :: String -> Plotter ()
xlabel = command . XLabel

-- | Y axis label
ylabel :: String -> Plotter ()
ylabel = command . YLabel

-- | Z axis label
zlabel :: String -> Plotter ()
zlabel = command . ZLabel

-- | plot title
title :: String -> Plotter ()
title = command . Title

-- | plot legend
legend :: [String] -> Plotter ()
legend = command . Legend

-- | Turn [Double] into a comma seperated list
-- of doubles formatted with 'realFloat'.
fromArray :: [Double] -> Builder
fromArray v = singleton '[' `mappend` fromArray' v
  where
    -- close the end
    fromArray' [] = singleton ']'
    -- if there is at least one, use a comma
    fromArray' (x0:xs@(_:_)) = realFloat x0 `mappend` fromString ", " `mappend` fromArray' xs
    -- if there is only one left, don't use a comma
    fromArray' (x:xs) = realFloat x `mappend` fromArray' xs

catBuilders :: [Builder] -> Builder
catBuilders (x:xs) = x `mappend` catBuilders xs
catBuilders [] = mempty

-- | export the plotter as text in matlab format
runMatlab :: Plotter a -> Text
runMatlab p = withPlotter p toMatlab
  where
    toOpts :: [String] -> String
    toOpts = concatMap (\o -> ", '" ++ o ++ "'")
    
    toMatlab :: PlotCommand -> Builder
    toMatlab (Figure Nothing) = fromString "figure()\n"
    toMatlab (Figure (Just k)) = fromString $ "figure(" ++ show k ++ ")\n"
    toMatlab (Plot xys opts) =
      catBuilders
      [ fromString "plot("
      , fromArray xs
      , fromString ", "
      , fromArray ys
      , fromString (toOpts opts ++")\n")
      ]
      where
        (xs, ys) = unzip xys
    toMatlab (Plot3 xyzs opts) =
      fromString $ "plot3(" ++ show xs ++ ", " ++ show ys ++ ", " ++ show zs
      ++ toOpts opts ++")\n"
      where
        (xs, ys, zs) = unzip3 xyzs
    toMatlab (Hold True) = fromString "hold on\n"
    toMatlab (Hold False) = fromString "hold off\n"
    toMatlab (Subplot kx ky kz) =
      fromString $ printf "subplot(%d, %d, %d)\n" kx ky kz
    toMatlab (XLabel name) = fromString ("xlabel('" ++ name ++ "')\n")
    toMatlab (YLabel name) = fromString ("ylabel('" ++ name ++ "')\n")
    toMatlab (ZLabel name) = fromString ("zlabel('" ++ name ++ "')\n")
    toMatlab (Title name) = fromString ("title('" ++ name ++ "')\n")
    toMatlab (Legend names) =
      fromString ("legend(" ++ intercalate ", " names' ++ ")\n")
      where
        names' = map (\n -> "'" ++ n ++ "'") names

---- | export the plotter as text in matplotlib format
--runMatplotlib :: Plotter a -> Text
--runMatplotlib = undefined

withPlotter :: Plotter a -> (PlotCommand -> Builder) -> Text
withPlotter (Plotter p) toBuilder =
  toLazyText $ catBuilders (map toBuilder (F.toList commands))
  where
    PlotterState commands = State.execState p (PlotterState mempty)