Source

textplot.next / TextPlot.hs

Full commit
{-# LANGUAGE TypeSynonymInstances #-}
--
-- Copyright (c) William Tennien Murphy 2011
--

module TextPlot
    ( -- * Types
      Function
    , Range
    , XYPlot(..), emptyXYPlot
    , ParamXYPlot(..), ParamFunction(..), emptyParamXYPlot
    , PolarPlot(..), PolarFunction(..), emptyPolarPlot
    -- * Ploting (how to plot)
    , PlotConfig(..), defaultConfig, Plot()
    -- * Output
    , plot
    , plotWithConfig
    , printPlot
    ) where

import Control.Monad (forM_, when)
import Control.Monad.ST (ST)
import Data.Array
import Data.Array.ST

type TextPlot = [[Char]]

-- | Range of values @(from, to)@.
type Range = (Double, Double)

-- some reasonable default values
defaultXrange :: Range
defaultXrange = (0.0, 1.0)
defaultYrange :: Range
defaultYrange = (0.0, 1.0)

-- | A function of one variable (@x -> y@).
type Function = Double -> Double
instance Show Function where show = const "<function::Double->Double>"

-- | Plot one or more functions (@x -> y@) in Cartesian coordinates.
data XYPlot = XYPlot {
      fp'xlim :: Range -- ^ limits of the abscissa (x) axis
    , fp'ylim :: Range -- ^ limits of the ordinate (y) axis
    , fp'functions :: [Function] -- ^ functions to plot
    } deriving Show

-- | A default empty 'XYPlot' with bounds of a unit square.
emptyXYPlot :: XYPlot
emptyXYPlot = XYPlot defaultXrange defaultYrange []

data ParamFunction = ParamFunction {
      xfun :: Function  -- ^ the first parametrized variable (@t -> x@)
    , yfun :: Function  -- ^ the second parametrized variable (@t -> y@)
    , tlim :: Range  -- ^ range of the free parameter @t@
    } deriving Show

-- | Plot one or more parametric functions in Cartesian coordiantes.
data ParamXYPlot = ParamXYPlot {
      param'xlim :: Range -- ^ limits of the abscissa (x) axis
    , param'ylim :: Range -- ^ limits of the ordinate (y) axis
    , param'functions :: [ParamFunction] -- ^ functions to plot
    } deriving Show

-- | A default empty 'ParamXYPlot'
emptyParamXYPlot :: ParamXYPlot
emptyParamXYPlot = ParamXYPlot defaultXrange defaultYrange []

data PolarFunction = PolarFunction {
      rfun :: Function -- ^ radius as a function of angle @phi@ (@phi -> r@)
    , philim :: (Double,Double) -- ^ range of the angle argument @phi@
    }  deriving Show

-- | Plot one or more functions in polar coordinates.
data PolarPlot = PolarPlot {
      polar'xlim :: Range -- ^ limits of the abscissa (x) axis
    , polar'ylim :: Range -- ^ limits of the ordinate (y) axis
    , polar'functions :: [PolarFunction] -- ^ functions to plot
    } deriving Show

-- | A default empty 'PolarPlot'
emptyPolarPlot :: PolarPlot
emptyPolarPlot = PolarPlot defaultXrange defaultYrange []

-- | Anything that can be plotted.
class Plot a where
    draw :: PlotConfig -> a -> TextPlot

data PlotConfig = PlotConfig {
      c'width :: Int     -- ^ plot width in characters
    , c'height :: Int    -- ^ plot height in characters
    , c'samples :: Int   -- ^ samples per line
    , c'showAxes :: Bool -- ^ draw axes or not (TODO)
    } deriving (Show, Eq)

-- | Default plot dimensions, suitable for 80x25 terminals.
defaultConfig :: PlotConfig
defaultConfig = PlotConfig 60 21 256 True

instance Plot XYPlot where
    draw (PlotConfig width height _ _) plt = fromArray $ runSTArray $ do
      arr <- createArray width height
      let xrange@(xmin,xmax) = fp'xlim plt
      let yrange = fp'ylim plt
      let dx = (xmax-xmin)/(fromIntegral width - 1)
      let xs = [ xmin + (fromIntegral c)*dx | c <- [0..width-1] ]
      forM_ (reverse (zip (fp'functions plt) symbols)) $
                \(f, sym) -> markPoints xrange yrange arr sym xs (map f xs)
      return arr

instance Plot ParamXYPlot where
    draw (PlotConfig width height samples _) plt = fromArray $ runSTArray $ do
      arr <- createArray width height
      let xrange = param'xlim plt
      let yrange = param'ylim plt
      let fns = param'functions plt
      forM_ (reverse (zip fns symbols)) $
            \(f, sym) -> do
              let (tmin,tmax) = tlim f
              let dt = (tmax-tmin)/(fromIntegral samples - 1)
              let ts = [ (fromIntegral t)*dt | t <- [0..samples-1] ]
              let xs = map (xfun f) ts
              let ys = map (yfun f) ts
              markPoints xrange yrange arr sym xs ys
      return arr

instance Plot PolarPlot where
    draw (PlotConfig width height samples _) plt = fromArray $ runSTArray $ do
      arr <- createArray width height
      let xrange = polar'xlim plt
      let yrange = polar'ylim plt
      let fns = polar'functions plt
      forM_ (reverse (zip fns symbols)) $
            \(f, sym) -> do
              let (phimin, phimax) = philim f
              let dphi = (phimax-phimin)/(fromIntegral samples - 1)
              let phis = [ (fromIntegral t)*dphi | t <- [0..samples-1] ]
              let rs = map (rfun f) phis
              let toCartesian (r,phi) = (r*cos phi, r*sin phi)
              let (xs,ys) = unzip . map toCartesian $ zip rs phis
              markPoints xrange yrange arr sym xs ys
      return arr

-- | Convert a plot to a multiline 'String' with default configuration
plot :: Plot p => p -> String
plot = plotWithConfig defaultConfig

-- | Convert a plot to multiline 'String' with custom configuration
plotWithConfig :: Plot p => PlotConfig -> p -> String
plotWithConfig config = unlines . draw config

-- | Print a plot with default configuration
printPlot :: Plot p => p -> IO()
printPlot = putStr . plot


{--------------------- backend array operations ---------------------------}


-- | Symbols to use for different plots.
symbols :: String
symbols = cycle "ox+#*@-"

-- | Create an 'STArray' of given screen dimensions
createArray :: Int -> Int -> ST s (STArray s (Int,Int) Char)
createArray width height = do
  let screenDims = ((0,0),(height-1,width-1))
  newArray screenDims ' ' :: ST s (STArray s (Int,Int) Char)

-- | Mark (x,y) points in a two-dimensional array of 'Char'
markPoints :: Range    -- ^ @x@ range
           -> Range    -- ^ @y@ range
           -> STArray s (Int,Int) Char -- ^ an array we operate on
           -> Char     -- ^ mark symbol
           -> [Double] -- ^ @xs@
           -> [Double] -- ^ @ys@
           -> ST s (STArray s (Int,Int) Char)
markPoints (xmin,xmax) (ymin,ymax) arr sym xs ys = do
  ((rmin,cmin),(rmax,cmax)) <- getBounds arr
  let width = cmax-cmin+1
  let height = rmax-rmin+1
  let w = fromIntegral width
  let h = fromIntegral height
  let dx = (xmax-xmin)/(w-1)  -- larger dx, dy steps to guarantee that
  let dy = (ymax-ymin)/(h-1)  -- max values stay within plot bounds
  let cols = [ round$(x-xmin)/dx | x <- xs ]
  let rows = [ round$(h-1-(y-ymin)/dy) | y <- ys ]
  forM_ (zip cols rows) $ \(c, r) ->
      when (r >= rmin && r <= rmax && c >= cmin && c < cmax) $
           writeArray arr (r,c) sym
  return arr

-- | Convert a two dimensional array to a list of lists
fromArray :: Array (Int,Int) a -> [[a]]
fromArray arr = splitEvery width (elems arr)
  where
    splitEvery :: Int -> [a] -> [[a]]
    splitEvery _ [] = []
    splitEvery n xs = (take n xs) : splitEvery n (drop n xs)
    width :: Int
    width = let ((_,colmin),(_,colmax)) = bounds arr
            in  colmax - colmin + 1