-----------------------------------------------------------------------------
-- |
-- Module      :  Graphics.Rendering.Chart.Plot.Vectors
-- Copyright   :  (c) Anton Vorontsov <anton@enomsg.org> 2014
-- License     :  BSD-style (see chart/COPYRIGHT)
--
-- Vector plots
--
{-# LANGUAGE CPP #-}
{-# LANGUAGE TemplateHaskell #-}

module Graphics.Rendering.Chart.Plot.Vectors(
    PlotVectors(..),
    VectorStyle(..),
    plotVectorField,
    plot_vectors_mapf,
    plot_vectors_grid,
    plot_vectors_title,
    plot_vectors_style,
    plot_vectors_scale,
    plot_vectors_values,
    vector_line_style,
    vector_head_style,
) where

import Control.Lens
import Control.Monad
#if !MIN_VERSION_base(4,8,0)
import Control.Applicative
#endif
import Data.Tuple
import Data.Colour hiding (over)
import Data.Colour.Names
import Data.Default.Class
import Graphics.Rendering.Chart.Axis
import Graphics.Rendering.Chart.Drawing
import Graphics.Rendering.Chart.Geometry
import Graphics.Rendering.Chart.Plot.Types

data VectorStyle = VectorStyle
    { VectorStyle -> LineStyle
_vector_line_style :: LineStyle
    , VectorStyle -> PointStyle
_vector_head_style :: PointStyle
    }

$( makeLenses ''VectorStyle )

data PlotVectors x y = PlotVectors
    { forall x y. PlotVectors x y -> String
_plot_vectors_title        :: String
    , forall x y. PlotVectors x y -> VectorStyle
_plot_vectors_style        :: VectorStyle
    -- | Set to 1 (default) to normalize the length of vectors to a space
    --   between them (so that the vectors never overlap on the graph).
    --   Set to 0 to disable any scaling.
    --   Values in between 0 and 1 are also permitted to adjust scaling.
    , forall x y. PlotVectors x y -> Double
_plot_vectors_scale        :: Double
    -- | Provide a square-tiled regular grid.
    , forall x y. PlotVectors x y -> [(x, y)]
_plot_vectors_grid         :: [(x,y)]
    -- | Provide a vector field (R^2 -> R^2) function.
    , forall x y. PlotVectors x y -> (x, y) -> (x, y)
_plot_vectors_mapf         :: (x,y) -> (x,y)
    -- | Provide a prepared list of (start,vector) pairs.
    , forall x y. PlotVectors x y -> [((x, y), (x, y))]
_plot_vectors_values       :: [((x,y),(x,y))]
    }

$( makeLenses ''PlotVectors )

mapGrid :: (PlotValue y, PlotValue x)
        => [(x,y)] -> ((x,y) -> (x,y)) -> [((x,y),(x,y))]
mapGrid :: forall y x.
(PlotValue y, PlotValue x) =>
[(x, y)] -> ((x, y) -> (x, y)) -> [((x, y), (x, y))]
mapGrid [(x, y)]
grid (x, y) -> (x, y)
f = forall a b. [a] -> [b] -> [(a, b)]
zip [(x, y)]
grid ((x, y) -> (x, y)
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(x, y)]
grid)

plotVectorField :: (PlotValue x, PlotValue y) => PlotVectors x y -> Plot x y
plotVectorField :: forall x y.
(PlotValue x, PlotValue y) =>
PlotVectors x y -> Plot x y
plotVectorField PlotVectors x y
pv = Plot
    { _plot_render :: PointMapFn x y -> BackendProgram ()
_plot_render     = forall x y.
(PlotValue x, PlotValue y) =>
PlotVectors x y -> PointMapFn x y -> BackendProgram ()
renderPlotVectors PlotVectors x y
pv
    , _plot_legend :: [(String, Rect -> BackendProgram ())]
_plot_legend     = [(forall x y. PlotVectors x y -> String
_plot_vectors_title PlotVectors x y
pv, forall x y.
(PlotValue x, PlotValue y) =>
PlotVectors x y -> Rect -> BackendProgram ()
renderPlotLegendVectors PlotVectors x y
pv)]
    , _plot_all_points :: ([x], [y])
_plot_all_points = (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(x, y)]
pts, forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(x, y)]
pts)
    }
  where
    pvals :: [((x, y), (x, y))]
pvals = forall x y. PlotVectors x y -> [((x, y), (x, y))]
_plot_vectors_values PlotVectors x y
pv
    mvals :: [((x, y), (x, y))]
mvals = forall y x.
(PlotValue y, PlotValue x) =>
[(x, y)] -> ((x, y) -> (x, y)) -> [((x, y), (x, y))]
mapGrid (forall x y. PlotVectors x y -> [(x, y)]
_plot_vectors_grid PlotVectors x y
pv) (forall x y. PlotVectors x y -> (x, y) -> (x, y)
_plot_vectors_mapf PlotVectors x y
pv)
    pts :: [(x, y)]
pts = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\((x, y)
a,(x, y)
b) -> [(x, y)
a,(x, y)
b]) ([((x, y), (x, y))]
pvals forall a. [a] -> [a] -> [a]
++ [((x, y), (x, y))]
mvals)

renderPlotVectors :: (PlotValue x, PlotValue y)
                  => PlotVectors x y -> PointMapFn x y -> BackendProgram ()
renderPlotVectors :: forall x y.
(PlotValue x, PlotValue y) =>
PlotVectors x y -> PointMapFn x y -> BackendProgram ()
renderPlotVectors PlotVectors x y
pv PointMapFn x y
pmap = do
    let pvals :: [((x, y), (x, y))]
pvals   = forall x y. PlotVectors x y -> [((x, y), (x, y))]
_plot_vectors_values PlotVectors x y
pv
        mvals :: [((x, y), (x, y))]
mvals   = forall y x.
(PlotValue y, PlotValue x) =>
[(x, y)] -> ((x, y) -> (x, y)) -> [((x, y), (x, y))]
mapGrid (forall x y. PlotVectors x y -> [(x, y)]
_plot_vectors_grid PlotVectors x y
pv) (forall x y. PlotVectors x y -> (x, y) -> (x, y)
_plot_vectors_mapf PlotVectors x y
pv)
        trans :: [((x, y), (x, y))]
trans   = forall {a} {a} {a} {b} {a} {a}.
(PlotValue a, PlotValue a, PlotValue a, PlotValue b, PlotValue a,
 PlotValue a) =>
((a, a), (a, a)) -> ((a, a), (a, b))
translateToStart forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([((x, y), (x, y))]
pvals forall a. [a] -> [a] -> [a]
++ [((x, y), (x, y))]
mvals)
        pvecs :: [(Point, Point)]
pvecs   = forall a. (a -> Bool) -> [a] -> [a]
filter (\(Point, Point)
v -> (Point, Point) -> Double
vlen' (Point, Point)
v forall a. Ord a => a -> a -> Bool
> Double
0) forall a b. (a -> b) -> a -> b
$ forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over forall (r :: * -> * -> *) a b.
Bitraversable r =>
Traversal (r a a) (r b b) a b
both (forall x y. PointMapFn x y -> (x, y) -> Point
mapXY PointMapFn x y
pmap) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [((x, y), (x, y))]
trans
        mgrid :: [Point]
mgrid   = forall a. Int -> [a] -> [a]
take Int
2 forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Point, Point)]
pvecs
        maxLen :: Double
maxLen  = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum forall a b. (a -> b) -> a -> b
$ (Point, Point) -> Double
vlen' forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Point, Point)]
pvecs
        spacing :: Double
spacing = (forall a. [a] -> Int -> a
!!Int
1) forall a b. (a -> b) -> a -> b
$ (Vector -> Double
vlen forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Point -> Point -> Vector
psub [Point]
mgrid (forall a. [a] -> [a]
reverse [Point]
mgrid)) forall a. [a] -> [a] -> [a]
++ [Double
maxLen]
        sfactor :: Double
sfactor = Double
spacingforall a. Fractional a => a -> a -> a
/Double
maxLen                  -- Non-adjusted scale factor
        afactor :: Double
afactor = Double
sfactor forall a. Num a => a -> a -> a
+ (Double
1 forall a. Num a => a -> a -> a
- Double
sfactor)forall a. Num a => a -> a -> a
*(Double
1 forall a. Num a => a -> a -> a
- forall x y. PlotVectors x y -> Double
_plot_vectors_scale PlotVectors x y
pv)
        tails :: [(Point, Point)]
tails   = Double -> (Point, Point) -> (Point, Point)
pscale Double
afactor forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Point, Point)]
pvecs          -- Paths of arrows' tails
        angles :: [Double]
angles  = (Vector -> Double
vangle forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Point, Point) -> Vector
psub' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> (b, a)
swap) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Point, Point)]
pvecs -- Angles of the arrows
        centers :: [Point]
centers = forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Point, Point)]
tails                     -- Where to draw arrow heads
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Double -> (Point, Point) -> BackendProgram ()
drawTail Double
radius) [(Point, Point)]
tails
    forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (Double -> Point -> Double -> BackendProgram ()
drawArrowHead Double
radius) [Point]
centers [Double]
angles
  where
    psub' :: (Point, Point) -> Vector
psub' = forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Point -> Point -> Vector
psub
    vlen' :: (Point, Point) -> Double
vlen' = Vector -> Double
vlen forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Point, Point) -> Vector
psub'
    pvs :: VectorStyle
pvs = forall x y. PlotVectors x y -> VectorStyle
_plot_vectors_style PlotVectors x y
pv
    radius :: Double
radius = PointStyle -> Double
_point_radius forall a b. (a -> b) -> a -> b
$ VectorStyle -> PointStyle
_vector_head_style VectorStyle
pvs
    hs :: Double -> PointStyle
hs Double
angle = VectorStyle -> PointStyle
_vector_head_style VectorStyle
pvs forall a b. a -> (a -> b) -> b
& Lens' PointStyle PointShape
point_shape
                  forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (\(PointShapeArrowHead Double
a) -> Double -> PointShape
PointShapeArrowHead forall a b. (a -> b) -> a -> b
$ Double
aforall a. Num a => a -> a -> a
+Double
angle)
    translateToStart :: ((a, a), (a, a)) -> ((a, a), (a, b))
translateToStart (s :: (a, a)
s@(a
x,a
y),(a
vx,a
vy)) = ((a, a)
s,(forall {a} {a} {a}.
(PlotValue a, PlotValue a, PlotValue a) =>
a -> a -> a
tr a
x a
vx,forall {a} {a} {a}.
(PlotValue a, PlotValue a, PlotValue a) =>
a -> a -> a
tr a
y a
vy))
      where tr :: a -> a -> a
tr a
p a
t = forall a. PlotValue a => Double -> a
fromValue forall a b. (a -> b) -> a -> b
$ forall a. PlotValue a => a -> Double
toValue a
p forall a. Num a => a -> a -> a
+ forall a. PlotValue a => a -> Double
toValue a
t
    pscale :: Double -> (Point, Point) -> (Point, Point)
pscale Double
w v :: (Point, Point)
v@(Point
s,Point
_) = (Point
s,Vector -> Point -> Point
translateP (Double -> Vector -> Vector
vscale Double
w forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Point, Point) -> Vector
psub' forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> (b, a)
swap (Point, Point)
v) Point
s)
    drawTail :: Double -> (Point, Point) -> BackendProgram ()
drawTail Double
r (Point, Point)
v = forall a. LineStyle -> BackendProgram a -> BackendProgram a
withLineStyle (VectorStyle -> LineStyle
_vector_line_style VectorStyle
pvs) forall a b. (a -> b) -> a -> b
$
        [Point] -> BackendProgram ()
strokePointPath forall a b. (a -> b) -> a -> b
$ (forall s a. s -> Getting (Endo [a]) s a -> [a]
^..forall s t a b. Each s t a b => Traversal s t a b
each) (Point, Point)
v'
      where
        v' :: (Point, Point)
v'  = Double -> (Point, Point) -> (Point, Point)
pscale (Double
1forall a. Num a => a -> a -> a
-(Double
3forall a. Fractional a => a -> a -> a
/Double
2)forall a. Num a => a -> a -> a
*Double
rforall a. Fractional a => a -> a -> a
/Double
l) (Point, Point)
v
        l :: Double
l   = (Point, Point) -> Double
vlen' (Point, Point)
v
    drawArrowHead :: Double -> Point -> Double -> BackendProgram ()
drawArrowHead Double
r (Point Double
x Double
y) Double
theta =
        forall a. Point -> BackendProgram a -> BackendProgram a
withTranslation (Double -> Double -> Point
Point (-Double
rforall a. Num a => a -> a -> a
*forall a. Floating a => a -> a
cos Double
theta) (-Double
rforall a. Num a => a -> a -> a
*forall a. Floating a => a -> a
sin Double
theta))
                        (PointStyle -> Point -> BackendProgram ()
drawPoint (Double -> PointStyle
hs Double
theta) (Double -> Double -> Point
Point Double
x Double
y))

renderPlotLegendVectors :: (PlotValue x, PlotValue y)
                        => PlotVectors x y -> Rect -> BackendProgram ()
renderPlotLegendVectors :: forall x y.
(PlotValue x, PlotValue y) =>
PlotVectors x y -> Rect -> BackendProgram ()
renderPlotLegendVectors PlotVectors x y
pv (Rect Point
p1 Point
p2) = do
    let y :: Double
y = (Point -> Double
p_y Point
p1 forall a. Num a => a -> a -> a
+ Point -> Double
p_y Point
p2)forall a. Fractional a => a -> a -> a
/Double
2
        pv' :: PlotVectors x y
pv' = forall x y. Lens' (PlotVectors x y) [(x, y)]
plot_vectors_grid forall s t a b. ASetter s t a b -> b -> s -> t
.~ []
            forall a b. (a -> b) -> a -> b
$ forall x y. Lens' (PlotVectors x y) [((x, y), (x, y))]
plot_vectors_values forall s t a b. ASetter s t a b -> b -> s -> t
.~ [((forall a. PlotValue a => Double -> a
fromValue forall a b. (a -> b) -> a -> b
$ Point -> Double
p_x Point
p1, forall a. PlotValue a => Double -> a
fromValue Double
y),
                                       (forall a. PlotValue a => Double -> a
fromValue forall a b. (a -> b) -> a -> b
$ Point -> Double
p_x Point
p2, forall a. PlotValue a => Double -> a
fromValue Double
0))]
            forall a b. (a -> b) -> a -> b
$ PlotVectors x y
pv
    forall x y.
(PlotValue x, PlotValue y) =>
PlotVectors x y -> PointMapFn x y -> BackendProgram ()
renderPlotVectors PlotVectors x y
pv' forall {a} {a}.
(PlotValue a, PlotValue a) =>
(Limit a, Limit a) -> Point
pmap
  where
    pmap :: (Limit a, Limit a) -> Point
pmap (LValue a
x,LValue a
y) = Double -> Double -> Point
Point (forall a. PlotValue a => a -> Double
toValue a
x) (forall a. PlotValue a => a -> Double
toValue a
y)
    pmap (Limit a, Limit a)
_ = Double -> Double -> Point
Point Double
0 Double
0

instance Default VectorStyle where
  def :: VectorStyle
def = VectorStyle
    { _vector_line_style :: LineStyle
_vector_line_style = (Double -> AlphaColour Double -> LineStyle
solidLine Double
lw forall a b. (a -> b) -> a -> b
$ forall a. Num a => Colour a -> AlphaColour a
opaque forall a. (Ord a, Floating a) => Colour a
blue)
                              { _line_cap :: LineCap
_line_cap = LineCap
LineCapSquare }
    , _vector_head_style :: PointStyle
_vector_head_style = AlphaColour Double
-> AlphaColour Double
-> Double
-> Double
-> PointShape
-> PointStyle
PointStyle (forall a. Num a => Colour a -> AlphaColour a
opaque forall a. (Ord a, Floating a) => Colour a
red) forall a. Num a => AlphaColour a
transparent Double
lw (Double
2forall a. Num a => a -> a -> a
*Double
lw)
                                      (Double -> PointShape
PointShapeArrowHead Double
0)
    } where lw :: Double
lw = Double
2

instance Default (PlotVectors x y) where
  def :: PlotVectors x y
def = PlotVectors
    { _plot_vectors_title :: String
_plot_vectors_title        = String
""
    , _plot_vectors_style :: VectorStyle
_plot_vectors_style        = forall a. Default a => a
def
    , _plot_vectors_scale :: Double
_plot_vectors_scale        = Double
1
    , _plot_vectors_grid :: [(x, y)]
_plot_vectors_grid         = []
    , _plot_vectors_mapf :: (x, y) -> (x, y)
_plot_vectors_mapf         = forall a. a -> a
id
    , _plot_vectors_values :: [((x, y), (x, y))]
_plot_vectors_values       = []
    }