plots-0.1.1.2: src/Plots/Types/HeatMap.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module : Plots.Types.HeatMap
-- Copyright : (C) 2016 Christopher Chalmers
-- License : BSD-style (see the file LICENSE)
-- Maintainer : Christopher Chalmers
-- Stability : experimental
-- Portability : non-portable
--
-- A heat map is a graphical representation of data where the individual
-- values contained in a matrix are represented as colours.
--
-- <<diagrams/src_Plots_Types_HeatMap_heatMapIndexedExample.svg#diagram=heatMapIndexedExample&height=350>>
--
-- (see 'heatMapIndexed' example for code to make this plot)
--
----------------------------------------------------------------------------
module Plots.Types.HeatMap
( -- * Heat map
HeatMap
, heatMap
, heatMap'
, heatMapIndexed
, heatMapIndexed'
-- * Lenses
, HasHeatMap (..)
-- ** Rendering functions
, pathHeatRender
, pixelHeatRender
, pixelHeatRender'
-- * Heat matrix
, HeatMatrix
, heatImage
, hmPoints
, hmSize
-- * Low level construction
, mkHeatMap
, mkHeatMatrix
, mkHeatMatrix'
) where
import Control.Lens hiding (transform, ( # ))
import Control.Monad.ST
import Control.Monad.State
import qualified Data.Foldable as F
import Data.Typeable
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Unboxed as V
import Data.Word (Word8)
import Codec.Picture
import Diagrams.Coordinates.Isomorphic
import Diagrams.Prelude
import Plots.Axis
import Plots.Style
import Plots.Types
------------------------------------------------------------------------
-- Heatmap
------------------------------------------------------------------------
-- | 2D Array of 'Double's.
data HeatMatrix = HeatMatrix
{ hmSize :: {-# UNPACK #-} !(V2 Int)
-- ^ The size of heat matrix.
, _hmVector :: {-# UNPACK #-} !(V.Vector Double)
, hmBoundLower :: {-# UNPACK #-} !Double
, hmBoundUpper :: {-# UNPACK #-} !Double
}
-- | Construct a heat matrix from a size and a generating function.
mkHeatMatrix :: V2 Int -> (V2 Int -> Double) -> HeatMatrix
mkHeatMatrix s@(V2 x y) f = runST $ do
mv <- M.new (x*y)
let go !q !a !b !i !j
| j == y = do v <- V.unsafeFreeze mv
return (HeatMatrix s v a b)
| i == x = go q a b 0 (j + 1)
| otherwise = do let !d = f (V2 i j)
M.unsafeWrite mv q d
go (q + 1) (min a d) (max b d) (i + 1) j
go 0 (1/0) (-1/0) 0 0
{-# INLINE mkHeatMatrix #-}
min' :: Double -> Double -> Double
min' !a !b
| isNaN b = a
| isInfinite b = a
| otherwise = min a b
{-# INLINE min' #-}
max' :: Double -> Double -> Double
max' !a !b
| isNaN b = a
| isInfinite b = a
| otherwise = max a b
{-# INLINE max' #-}
data MM = MM {-# UNPACK #-} !Double {-# UNPACK #-} !Double
-- | Compute the minimum and maximum of a vector in one pass. This
-- ignores any @Infinity@ or @NaN@ values (since these make no sense
-- for colour bar ranges).
minMax :: V.Vector Double -> (Double, Double)
minMax = fini . V.foldl' go (MM (1/0) (-1/0))
where
go (MM lo hi) k = MM (min' lo k) (max' hi k)
fini (MM lo hi) = (lo, hi)
{-# INLINE minMax #-}
-- | Construct a heat matrix from a foldable of foldables.
--
-- @
-- 'mkHeatMatrix'' :: [['Double']] -> 'HeatMatrix'
-- 'mkHeatMatrix'' :: ['Vector' 'Double'] -> 'HeatMatrix'
-- @
mkHeatMatrix' :: (F.Foldable f, F.Foldable g) => f (g Double) -> HeatMatrix
mkHeatMatrix' xss = HeatMatrix (V2 x y) vd a b
where
(a,b) = minMax vd
vd = V.create $ do
mv <- M.new (x*y)
let go !_ [] = return mv
go j (r:rs) = V.unsafeCopy (M.unsafeSlice (j*x) x mv) r >> go (j-1) rs
go (y - 1) vs
-- vs is in reverse since we used foldl' to build it
(!x,!y,!vs) = F.foldl' f (maxBound,0,[]) xss
f (!i,!j,!ss) xs = let !v = V.fromList (F.toList xs)
in (min i (V.length v), j+1, v : ss)
-- | Indexed traversal over the values of a 'HeatMatrix'.
hmPoints :: IndexedTraversal' (V2 Int) HeatMatrix Double
hmPoints f (HeatMatrix e@(V2 x y) v a b) =
go 0 0 0 <&> \vs ->
let v'= V.fromListN (x*y) vs
in HeatMatrix e v' a b
where
-- V2 x y = hmExtent
go !s !i !j
| i >= x = go s 0 (j+1)
| j >= y = pure []
| otherwise = (:) <$> indexed f (V2 i j) (V.unsafeIndex v s)
<*> go (s+1) (i+1) j
{-# INLINE [0] hmPoints #-}
{-# RULES
"hmPoints/foldr"
hmPoints = ifoldring hmFold :: Getting (Endo r) HeatMatrix Double;
"hmPoints/ifoldr"
hmPoints = ifoldring hmFold :: IndexedGetting (V2 Int) (Endo r) HeatMatrix Double
#-}
hmFold :: (V2 Int -> Double -> b -> b) -> b -> HeatMatrix -> b
hmFold f b0 (HeatMatrix (V2 x y) v _ _) = go 0 0 0 b0 where
go !s !i !j b
| i >= x = go s 0 (j+1) b
| j >= y = b
| otherwise = f (V2 i j) (V.unsafeIndex v s) (go (s+1) (i+1) j b)
{-# INLINE hmFold #-}
-- Rendering heat matrices --------------------------------------------
-- | Render an heatmap as an 'ImageRGB8'.
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_HeatMap_pixelHeatRenderExample.svg#diagram=pixelHeatRenderExample&height=350>>
--
-- > import Plots
-- >
-- > pixelHeatRenderExample =
-- > let f (V2 x y) = fromIntegral x + fromIntegral y
-- > myHM = mkHeatMatrix (V2 5 5) f
-- > in pixelHeatRender myHM viridis
--
pixelHeatRender
:: (Renderable (DImage n Embedded) b, TypeableFloat n)
=> HeatMatrix
-> ColourMap
-> QDiagram b V2 n Any
pixelHeatRender hm cm =
alignBL . image $ DImage (ImageRaster (ImageRGB8 img)) x y mempty
where
img = heatImage hm cm
V2 x y = hmSize hm
-- | Render an heatmap as an 'ImageRGB8' with @n@ pixels per heat matrix
-- point.
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_HeatMap_pixelHeatRenderExample'.svg#diagram=pixelHeatRenderExample'&height=350>>
--
-- > import Plots
-- >
-- > pixelHeatRenderExample' =
-- > let f (V2 x y) = fromIntegral x + fromIntegral y
-- > myHM = mkHeatMatrix (V2 5 5) f
-- > in pixelHeatRender' 10 myHM viridis
--
pixelHeatRender'
:: (Renderable (DImage n Embedded) b, TypeableFloat n)
=> Int
-> HeatMatrix
-> ColourMap
-> QDiagram b V2 n Any
pixelHeatRender' n hm cm =
scale (1/fromIntegral n) . alignBL . image $ DImage (ImageRaster (ImageRGB8 img)) (x*n) (y*n) mempty
where
img = scaleImage n $ heatImage hm cm
V2 x y = hmSize hm
-- | Scale an image so each pixel takes (n*n) pixels. This can be
-- useful for using 'heatImage' on small heat matrices to give a
-- sharper image.
scaleImage :: Int -> Image PixelRGB8 -> Image PixelRGB8
scaleImage n img | n == 1 = img
| n == 0 = Image 0 0 S.empty
| n < 0 = error "scaleImage: negative scale"
scaleImage n (Image x y v) = Image (n*x) (n*y) vn where
!refV = V.fromList $ map (*3) [ i + n*x*j | i <- [0..n-1], j <- [0..n-1] ]
!n3 = 3*n
vn = S.create $ do
mv <- M.new (n * n * S.length v)
let go !q !i !s | q >= 3*x*y = return mv
| i == x = go q 0 (s + 3*x*n*(n-1))
go q i s = do
let !r = S.unsafeIndex v q
!g = S.unsafeIndex v (q+1)
!b = S.unsafeIndex v (q+2)
V.forM_ refV $ \ds -> do
M.unsafeWrite mv (s + ds ) r
M.unsafeWrite mv (s + ds + 1) g
M.unsafeWrite mv (s + ds + 2) b
go (q+3) (i+1) (s+n3)
go 0 0 0
-- | Create an image of 'PixelsRGB8' using the heat matrix.
heatImage :: HeatMatrix -> ColourMap -> Image PixelRGB8
heatImage (HeatMatrix (V2 x y) dv a b) cm = Image x y v where
!cv = mkColourVector cm
v :: S.Vector Word8
v = S.create $ do
mv <- M.new (3 * x * y)
let !m = 256 / (b - a)
writeColour q (toSRGB24 -> RGB rr gg bb) = do
M.unsafeWrite mv q rr
M.unsafeWrite mv (q+1) gg
M.unsafeWrite mv (q+2) bb
colourValue !q !d
| isNaN d = writeColour q (cm^.nanColour)
| isInfinite d = writeColour q (if d < 0 then cm^.negInfColour else cm^.infColour)
| otherwise = do
let !o = 3 * (min 255 . max 0 . round $ (d - a) * m)
M.unsafeWrite mv q (V.unsafeIndex cv o )
M.unsafeWrite mv (q+1) (V.unsafeIndex cv (o+1))
M.unsafeWrite mv (q+2) (V.unsafeIndex cv (o+2))
go s i q
| s == x-1 = do
let d = V.unsafeIndex dv s
colourValue q d
return mv
| i == x = go (s - 2*x) 0 q
| otherwise = do
let d = V.unsafeIndex dv s
colourValue q d
go (s+1) (i+1) (q+3)
go (x * (y-1)) 0 0
-- Make an unboxed colour map using 256 samples.
mkColourVector :: ColourMap -> V.Vector Word8
mkColourVector cm = V.create $ do
mv <- M.new (3*256)
let go i | i == 3*256 = return mv
| otherwise = do
let x = fromIntegral i / (3*256)
RGB r g b = cm ^. ixColourR x . to toSRGB24
M.unsafeWrite mv i r
M.unsafeWrite mv (i+1) g
M.unsafeWrite mv (i+2) b
go (i+3)
go 0
-- | Render the heat map as a collection squares made up of 'Trail's.
-- This method is compatible with all backends and should always look
-- sharp. However it can become slow and large for large heat maps.
--
-- It is recommended to use 'pathHeatRender' for small heat maps and
-- 'pixelHeatRender' for larger ones.
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_HeatMap_pathHeatRenderExample.svg#diagram=pathHeatRenderExample&height=350>>
--
-- > import Plots
-- >
-- > pathHeatRenderExample =
-- > let f (V2 x y) = fromIntegral x + fromIntegral y
-- > myHM = mkHeatMatrix (V2 5 5) f
-- > in pathHeatRender myHM viridis
--
pathHeatRender
:: (Renderable (Path V2 n) b, TypeableFloat n)
=> HeatMatrix
-> ColourMap
-> QDiagram b V2 n Any
pathHeatRender hm@(HeatMatrix _ _ a b) cm = ifoldMapOf hmPoints mk hm # lwO 0
where
normalise d = (d - a) / (b - a)
mk v@(V2 i j) d =
rect w h
# alignTR
# translate (fromIntegral <$> v ^+^ 1)
# fc (cm ^. ixColour (normalise d))
where
-- Squares that are not on the top left edge are slightly
-- bigger to remove phantom gaps
w | i == 0 = 1
| otherwise = 1.5
h | j == 0 = 1
| otherwise = 1.5
------------------------------------------------------------------------
-- Heat matrix
------------------------------------------------------------------------
-- | A mapping from points in a 2D axis do 'Double's. These 'Double's
-- are converted to colours using the axis 'ColourMap'.
data HeatMap b n = HeatMap
{ hMatrix :: HeatMatrix
, hStart :: P2 n
, hSize :: V2 n
, hGridSty :: Style V2 n
, hGridVisible :: Bool
, hLimits :: Maybe (Double,Double)
, hDraw :: HeatMatrix -> ColourMap -> QDiagram b V2 n Any
} deriving Typeable
type instance V (HeatMap b n) = V2
type instance N (HeatMap b n) = n
-- | Class of things that let you change the heatmap options.
class HasHeatMap f a b | a -> b where
-- | Lens onto the heatmap options.
heatMapOptions :: LensLike' f a (HeatMap b (N a))
-- | Whether there should be grid lines draw for the heat map.
--
-- Default is 'False'.
heatMapGridVisible :: Functor f => LensLike' f a Bool
heatMapGridVisible = heatMapOptions . lens hGridVisible (\s b -> (s {hGridVisible = b}))
-- | The style applied to the grid lines for the heat map, if they're
-- visible.
--
-- Default is 'mempty'.
heatMapGridStyle :: Functor f => LensLike' f a (Style V2 (N a))
heatMapGridStyle = heatMapOptions . lens hGridSty (\s b -> (s {hGridSty = b}))
-- | The size of each individual square in the heat map.
--
-- Default is @'V2' 1 1@.
heatMapSize :: Functor f => LensLike' f a (V2 (N a))
heatMapSize = heatMapOptions . lens hSize (\s b -> (s {hSize = b}))
-- | The size of the full extent of the heat map.
--
-- Default is extent of the heat matrix.
heatMapExtent :: (Functor f, Fractional (N a)) => LensLike' f a (V2 (N a))
heatMapExtent = heatMapOptions . l where
l f hm = f (hSize hm * s) <&> \x -> hm { hSize = x / s }
where s = fmap fromIntegral (hmSize $ hMatrix hm)
-- | The starting point at the bottom left corner of the heat map.
--
-- Default is 'origin'
heatMapStart :: Functor f => LensLike' f a (P2 (N a))
heatMapStart = heatMapOptions . lens hStart (\s b -> (s {hStart = b}))
-- | The center point of the heat map.
heatMapCentre :: (Functor f, Fractional (N a)) => LensLike' f a (P2 (N a))
heatMapCentre = heatMapOptions . l where
l f hm = f (hStart hm .+^ v) <&> \p -> hm { hStart = p .-^ v }
where v = fmap fromIntegral (hmSize $ hMatrix hm) * hSize hm / 2
-- | Limits @(a,b)@ used on the data such that @a@ is the start of the
-- 'ColourMap' and @b@ is the end of the 'ColourMap'. Default is @(0,1)@.
heatMapLimits :: Functor f => LensLike' f a (Maybe (Double, Double))
heatMapLimits = heatMapOptions . lens hLimits (\s b -> (s {hLimits = b}))
-- | Funtion used to render the heat map. See 'pathHeatRender' and
-- 'pixelHeatRender'.
--
-- Default is 'pathHeatRender'.
heatMapRender :: Functor f => LensLike' f a (HeatMatrix -> ColourMap -> QDiagram b V2 (N a) Any)
heatMapRender = heatMapOptions . lens hDraw (\s b -> (s {hDraw = b}))
instance HasHeatMap f (HeatMap b n) b where
heatMapOptions = id
instance (Functor f, HasHeatMap f a b) => HasHeatMap f (Plot a b) b where
heatMapOptions = rawPlot . heatMapOptions
instance OrderedField n => Enveloped (HeatMap b n) where
getEnvelope hm = getEnvelope (fromCorners p (p .+^ v))
where p = view heatMapStart hm
v = view heatMapExtent hm
instance (Typeable b, TypeableFloat n, Renderable (Path V2 n) b)
=> Plotable (HeatMap b n) b where
renderPlotable s _sty HeatMap {..} =
transform (s^.specTrans) $
grid <> hDraw matrix' (s^.specColourMap)
# transform (scaleV hSize)
# moveTo hStart
where
--- TODO
grid = mempty
--- XXX need to give _range to the axis somehow (for colour bar range)
matrix' = case hLimits of
-- Just r@(a,b) -> (r, hMatrix { hmFun = (/ (b - a)) . (+a) . hmFun hMatrix })
-- Nothing -> normaliseHeatMatrix hMatrix
Just (a,b) -> hMatrix { hmBoundLower = a, hmBoundUpper = b }
Nothing -> hMatrix
-- XXX make better
defLegendPic sty HeatMap {..} = square 5 # applyAreaStyle sty
scaleV :: (Additive v, Fractional n) => v n -> Transformation v n
scaleV v = fromLinear f f
where f = (liftU2 (*) v) <-> (\u -> liftU2 (/) u v)
-- | Construct a 'Heatmap' using the given 'HeatMatrix'.
mkHeatMap :: (Renderable (Path V2 n) b, TypeableFloat n)
=> HeatMatrix -> HeatMap b n
mkHeatMap mat = HeatMap
{ hMatrix = mat
, hStart = origin
, hSize = V2 1 1
, hGridSty = mempty
, hGridVisible = False
, hLimits = Nothing
, hDraw = pathHeatRender
}
-- Adding to axis ------------------------------------------------------
-- | Add a 'HeatMap' plot using the extent of the heatmap and a
-- generating function.
--
-- @
-- 'heatMap' :: [['Double']] -> 'State' ('Plot' ('HeatMap' b n)) () -> 'State' ('Axis' b 'V2' n) ()
-- @
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_HeatMap_heatMapExample.svg#diagram=heatMapExample&height=350>>
--
-- > import Plots
-- > heatMapAxis :: Axis B V2 Double
-- > heatMapAxis = r2Axis &~ do
-- > display colourBar
-- > axisExtend .= noExtend
-- >
-- > let xs = [[1,2,3],[4,5,6]]
-- > heatMap xs $ heatMapSize .= V2 10 10
--
-- > heatMapExample = renderAxis heatMapAxis
--
heatMap
:: (F.Foldable f,
F.Foldable g,
TypeableFloat n,
Typeable b,
MonadState (Axis b V2 n) m,
Renderable (Path V2 n) b)
=> f (g Double)
-> State (Plot (HeatMap b n) b) ()
-- ^ changes to plot options
-> m () -- ^ add plot to 'Axis'
heatMap xss s = do
let hm@(HeatMatrix _ _ a b) = mkHeatMatrix' xss
addPlotable (mkHeatMap hm) s
-- (don't like this way of doing it)
colourBarRange .= over both realToFrac (a,b)
-- | Add a 'HeatMap' plot using the extent of the heatmap and a
-- generating function.
--
-- @
-- 'heatMap'' :: [['Double']] -> 'State' ('Axis' b 'V2' n) ()
-- @
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_HeatMap_heatMapExample'.svg#diagram=heatMapExample'&height=350>>
--
-- > import Plots
-- > heatMapAxis' :: Axis B V2 Double
-- > heatMapAxis' = r2Axis &~ do
-- > display colourBar
-- > axisExtend .= noExtend
-- > axisColourMap .= Plots.magma
-- >
-- > let xs = [[1,2,3],[4,5,6]]
-- > heatMap' xs
--
-- > heatMapExample' = renderAxis heatMapAxis'
--
heatMap'
:: (F.Foldable f,
F.Foldable g,
TypeableFloat n,
Typeable b,
MonadState (Axis b V2 n) m,
Renderable (Path V2 n) b)
=> f (g Double)
-> m () -- ^ add plot to 'Axis'
heatMap' xss = heatMap xss (return ())
-- | Add a 'HeatMap' plot using the extent of the heatmap and a
-- generating function.
--
-- @
-- 'heatMapIndexed' :: 'V2' 'Int' -> ('V2' 'Int' -> 'Double') -> 'State' ('Plot' ('HeatMap' b n)) () -> 'State' ('Axis' b 'V2' n) ()
-- 'heatMapIndexed' :: ('Int', 'Int') -> (('Int', 'Int') -> 'Double') -> 'State' ('Plot' ('HeatMap' b n)) () -> 'State' ('Axis' b 'V2' n) ()
-- @
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_HeatMap_heatMapIndexedExample.svg#diagram=heatMapIndexedExample&height=350>>
--
-- > import Plots
-- > heatMapIndexedAxis :: Axis B V2 Double
-- > heatMapIndexedAxis = r2Axis &~ do
-- > display colourBar
-- > axisExtend .= noExtend
-- >
-- > let f (V2 x y) = fromIntegral x + fromIntegral y
-- > heatMapIndexed (V2 3 3) f $ heatMapSize .= V2 10 10
--
-- > heatMapIndexedExample = renderAxis heatMapIndexedAxis
--
heatMapIndexed
:: (VectorLike V2 Int i,
TypeableFloat n,
Typeable b,
MonadState (Axis b V2 n) m,
Renderable (Path V2 n) b)
=> i -- ^ extent of array
-> (i -> Double) -- ^ heat from index
-> State (Plot (HeatMap b n) b) ()
-- ^ changes to plot options
-> m () -- ^ add plot to 'Axis'
heatMapIndexed i f s = do
let hm@(HeatMatrix _ _ a b) = mkHeatMatrix (view unvectorLike i) (f . view vectorLike)
addPlotable (mkHeatMap hm) s
-- (don't like this way of doing it)
colourBarRange .= over both realToFrac (a,b)
-- | Add a 'HeatMap' plot using the extent of the heatmap and a
-- generating function without changes to the heap map options.
--
-- @
-- 'heatMapIndexed' :: 'V2' 'Int' -> ('V2' 'Int' -> 'Double') -> 'State' ('Axis' b 'V2' n) ()
-- 'heatMapIndexed' :: ('Int', 'Int') -> (('Int', 'Int') -> 'Double') -> 'State' ('Axis' b 'V2' n) ()
-- @
--
-- === __Example__
--
-- <<diagrams/src_Plots_Types_HeatMap_heatMapIndexedExample'.svg#diagram=heatMapIndexedExample'&height=350>>
--
-- > import Plots
-- > heatMapIndexedAxis' :: Axis B V2 Double
-- > heatMapIndexedAxis' = r2Axis &~ do
-- > display colourBar
-- > axisExtend .= noExtend
-- > axisColourMap .= Plots.magma
-- >
-- > let f (V2 x y) = fromIntegral x + fromIntegral y
-- > heatMapIndexed' (V2 3 3) f
--
-- > heatMapIndexedExample' = renderAxis heatMapIndexedAxis'
--
heatMapIndexed'
:: (VectorLike V2 Int i,
TypeableFloat n,
Typeable b,
MonadState (Axis b V2 n) m,
Renderable (Path V2 n) b)
=> i -- ^ extent of array
-> (i -> Double) -- ^ heat from index
-> m () -- ^ add plot to 'Axis'
heatMapIndexed' i f = heatMapIndexed i f (return ())