{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeFamilies        #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Physics.ForceLayout
-- Copyright   :  (c) 2011 Brent Yorgey
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  byorgey@cis.upenn.edu
--
-- A simple, Haskell-native simulator for doing force-directed layout,
-- /e.g./ of trees or graphs.
--
-- To use, just create an 'Ensemble' like so:
--
-- > import           Control.Lens        ((&), (.~))
-- > import           Data.Default        (def)
-- > import qualified Data.Map            as M
-- > import           Linear.Affine
-- > import           Linear.V2
-- > import           Physics.ForceLayout
-- >
-- > e :: Ensemble V2 Double
-- > e = Ensemble [ (edges,    hookeForce 0.05 4)
-- >              , (allPairs, coulombForce 1)
-- >              ]
-- >              particleMap
-- >   where edges       = [(1,2), (2,3), (2,5), (3,5), (3,4), (4,5)]
-- >         allPairs    = [(x,y) | x <- [1..4], y <- [x+1..5]]
-- >         particleMap = M.fromList . zip [1..]
-- >                     . map (initParticle . P . uncurry V2)
-- >                     $ [ (2.0, 3.1), (6.3, 7.2)
-- >                       , (0.3, 4.2), (1.6, -1.1)
-- >                       , (4.8, 2.9)
-- >                       ]
--
-- Then run a simulation using either 'simulate' (to get the list of
-- all intermediate states) or 'forceLayout' (to get only the ending
-- state):
--
-- > e' :: Ensemble V2 Double
-- > e' = forceLayout (def & damping     .~ 0.8
-- >                       & energyLimit .~ Just 0.001
-- >                       & stepLimit   .~ Nothing
-- >                  )
-- >                  e
--
-- See the diagrams-contrib package
-- (<http://github.com/diagrams/diagrams-contrib/>) for more
-- examples.
-----------------------------------------------------------------------------

module Physics.ForceLayout
       ( -- * Data structures

         Particle(..), pos, vel, force
       , initParticle

       , PID
       , Edge
       , Ensemble(..), forces, particles

         -- * Pre-defined forces

       , hookeForce
       , coulombForce
       , distForce

         -- * Running simulations

       , ForceLayoutOpts(..)
       , damping, energyLimit, stepLimit
       , simulate
       , forceLayout

         -- * Internals

       , ensembleStep
       , particleStep
       , recalcForces
       , kineticEnergy

       ) where

import           Data.Default
import qualified Data.Foldable      as F
import qualified Data.Map           as M
import           Data.Monoid

import           Control.Lens

import           Linear.Affine
import           Linear.Metric
import           Linear.Vector

------------------------------------------------------------
--  Particles
------------------------------------------------------------

-- | A particle has a current position, current velocity, and current
--   force acting on it.
data Particle v n = Particle { forall (v :: * -> *) n. Particle v n -> Point v n
_pos   :: Point v n
                             , forall (v :: * -> *) n. Particle v n -> v n
_vel   :: v n
                             , forall (v :: * -> *) n. Particle v n -> v n
_force :: v n
                             }
  deriving Particle v n -> Particle v n -> Bool
(Particle v n -> Particle v n -> Bool)
-> (Particle v n -> Particle v n -> Bool) -> Eq (Particle v n)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (v :: * -> *) n.
Eq (v n) =>
Particle v n -> Particle v n -> Bool
$c== :: forall (v :: * -> *) n.
Eq (v n) =>
Particle v n -> Particle v n -> Bool
== :: Particle v n -> Particle v n -> Bool
$c/= :: forall (v :: * -> *) n.
Eq (v n) =>
Particle v n -> Particle v n -> Bool
/= :: Particle v n -> Particle v n -> Bool
Eq

makeLenses ''Particle

-- | Create an initial particle at rest at a particular location.
initParticle :: (Additive v, Num n) => Point v n -> Particle v n
initParticle :: forall (v :: * -> *) n.
(Additive v, Num n) =>
Point v n -> Particle v n
initParticle Point v n
p = Point v n -> v n -> v n -> Particle v n
forall (v :: * -> *) n. Point v n -> v n -> v n -> Particle v n
Particle Point v n
p v n
forall a. Num a => v a
forall (f :: * -> *) a. (Additive f, Num a) => f a
zero v n
forall a. Num a => v a
forall (f :: * -> *) a. (Additive f, Num a) => f a
zero

------------------------------------------------------------
--  Ensembles
------------------------------------------------------------

-- | Used to uniquely identify particles.
type PID = Int

-- | An edge is a pair of particle identifiers.
type Edge = (PID, PID)

-- | An @Ensemble@ is a physical configuration of particles.  It
--   consists of a mapping from particle IDs (unique integers) to
--   particles, and a list of forces that are operative.  Each force
--   has a list of edges to which it applies, and is represented by a
--   function giving the force between any two points.
data Ensemble v n = Ensemble { forall (v :: * -> *) n.
Ensemble v n -> [([Edge], Point v n -> Point v n -> v n)]
_forces    :: [([Edge], Point v n -> Point v n -> v n)]
                             , forall (v :: * -> *) n. Ensemble v n -> Map PID (Particle v n)
_particles :: M.Map PID (Particle v n)
                             }

makeLenses ''Ensemble

------------------------------------------------------------
--  Simulation internals
------------------------------------------------------------

-- | Simulate one time step for an entire ensemble, with the given
--   damping factor.
ensembleStep :: (Additive v, Num n) => n -> Ensemble v n -> Ensemble v n
ensembleStep :: forall (v :: * -> *) n.
(Additive v, Num n) =>
n -> Ensemble v n -> Ensemble v n
ensembleStep n
d = (ASetter
  (Ensemble v n)
  (Ensemble v n)
  (Map PID (Particle v n))
  (Map PID (Particle v n))
-> (Map PID (Particle v n) -> Map PID (Particle v n))
-> Ensemble v n
-> Ensemble v n
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter
  (Ensemble v n)
  (Ensemble v n)
  (Map PID (Particle v n))
  (Map PID (Particle v n))
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(Map PID (Particle v n) -> f (Map PID (Particle v n)))
-> Ensemble v n -> f (Ensemble v n)
particles ((Map PID (Particle v n) -> Map PID (Particle v n))
 -> Ensemble v n -> Ensemble v n)
-> ((Particle v n -> Particle v n)
    -> Map PID (Particle v n) -> Map PID (Particle v n))
-> (Particle v n -> Particle v n)
-> Ensemble v n
-> Ensemble v n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Particle v n -> Particle v n)
-> Map PID (Particle v n) -> Map PID (Particle v n)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map) (n -> Particle v n -> Particle v n
forall (v :: * -> *) n.
(Additive v, Num n) =>
n -> Particle v n -> Particle v n
particleStep n
d) (Ensemble v n -> Ensemble v n)
-> (Ensemble v n -> Ensemble v n) -> Ensemble v n -> Ensemble v n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ensemble v n -> Ensemble v n
forall (v :: * -> *) n.
(Additive v, Num n) =>
Ensemble v n -> Ensemble v n
recalcForces

-- | Simulate one time step for a particle (assuming the force acting
--   on it has already been computed), with the given damping factor.
particleStep :: (Additive v, Num n) => n -> Particle v n -> Particle v n
particleStep :: forall (v :: * -> *) n.
(Additive v, Num n) =>
n -> Particle v n -> Particle v n
particleStep n
d = Particle v n -> Particle v n
forall {v :: * -> *} {n}.
(Additive v, Num n) =>
Particle v n -> Particle v n
stepPos (Particle v n -> Particle v n)
-> (Particle v n -> Particle v n) -> Particle v n -> Particle v n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Particle v n -> Particle v n
stepVel
  where stepVel :: Particle v n -> Particle v n
stepVel Particle v n
p = (v n -> Identity (v n)) -> Particle v n -> Identity (Particle v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(v n -> f (v n)) -> Particle v n -> f (Particle v n)
vel ((v n -> Identity (v n))
 -> Particle v n -> Identity (Particle v n))
-> v n -> Particle v n -> Particle v n
forall s t a b. ASetter s t a b -> b -> s -> t
.~ (n
d n -> v n -> v n
forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ (Particle v n
pParticle v n -> Getting (v n) (Particle v n) (v n) -> v n
forall s a. s -> Getting a s a -> a
^.Getting (v n) (Particle v n) (v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(v n -> f (v n)) -> Particle v n -> f (Particle v n)
vel v n -> v n -> v n
forall a. Num a => v a -> v a -> v a
forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a
^+^ Particle v n
pParticle v n -> Getting (v n) (Particle v n) (v n) -> v n
forall s a. s -> Getting a s a -> a
^.Getting (v n) (Particle v n) (v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(v n -> f (v n)) -> Particle v n -> f (Particle v n)
force)) (Particle v n -> Particle v n) -> Particle v n -> Particle v n
forall a b. (a -> b) -> a -> b
$ Particle v n
p
        stepPos :: Particle v n -> Particle v n
stepPos Particle v n
p = (Point v n -> Identity (Point v n))
-> Particle v n -> Identity (Particle v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(Point v n -> f (Point v n)) -> Particle v n -> f (Particle v n)
pos ((Point v n -> Identity (Point v n))
 -> Particle v n -> Identity (Particle v n))
-> (Point v n -> Point v n) -> Particle v n -> Particle v n
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (Point v n -> Diff (Point v) n -> Point v n
forall a. Num a => Point v a -> Diff (Point v) a -> Point v a
forall (p :: * -> *) a. (Affine p, Num a) => p a -> Diff p a -> p a
.+^ Particle v n
pParticle v n -> Getting (v n) (Particle v n) (v n) -> v n
forall s a. s -> Getting a s a -> a
^.Getting (v n) (Particle v n) (v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(v n -> f (v n)) -> Particle v n -> f (Particle v n)
vel) (Particle v n -> Particle v n) -> Particle v n -> Particle v n
forall a b. (a -> b) -> a -> b
$ Particle v n
p

-- | Recalculate all the forces acting in the next time step of an
--   ensemble.
recalcForces :: (Additive v, Num n) => Ensemble v n -> Ensemble v n
recalcForces :: forall (v :: * -> *) n.
(Additive v, Num n) =>
Ensemble v n -> Ensemble v n
recalcForces = Ensemble v n -> Ensemble v n
forall (v :: * -> *) n.
(Additive v, Num n) =>
Ensemble v n -> Ensemble v n
calcForces (Ensemble v n -> Ensemble v n)
-> (Ensemble v n -> Ensemble v n) -> Ensemble v n -> Ensemble v n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ensemble v n -> Ensemble v n
zeroForces
  where zeroForces :: Ensemble v n -> Ensemble v n
zeroForces = ((Map PID (Particle v n) -> Identity (Map PID (Particle v n)))
-> Ensemble v n -> Identity (Ensemble v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(Map PID (Particle v n) -> f (Map PID (Particle v n)))
-> Ensemble v n -> f (Ensemble v n)
particles ((Map PID (Particle v n) -> Identity (Map PID (Particle v n)))
 -> Ensemble v n -> Identity (Ensemble v n))
-> (Map PID (Particle v n) -> Map PID (Particle v n))
-> Ensemble v n
-> Ensemble v n
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~) ((Map PID (Particle v n) -> Map PID (Particle v n))
 -> Ensemble v n -> Ensemble v n)
-> ((Particle v n -> Particle v n)
    -> Map PID (Particle v n) -> Map PID (Particle v n))
-> (Particle v n -> Particle v n)
-> Ensemble v n
-> Ensemble v n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Particle v n -> Particle v n)
-> Map PID (Particle v n) -> Map PID (Particle v n)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((Particle v n -> Particle v n) -> Ensemble v n -> Ensemble v n)
-> (Particle v n -> Particle v n) -> Ensemble v n -> Ensemble v n
forall a b. (a -> b) -> a -> b
$ (v n -> Identity (v n)) -> Particle v n -> Identity (Particle v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(v n -> f (v n)) -> Particle v n -> f (Particle v n)
force ((v n -> Identity (v n))
 -> Particle v n -> Identity (Particle v n))
-> v n -> Particle v n -> Particle v n
forall s t a b. ASetter s t a b -> b -> s -> t
.~ v n
forall a. Num a => v a
forall (f :: * -> *) a. (Additive f, Num a) => f a
zero
        calcForces :: Ensemble v n -> Ensemble v n
calcForces (Ensemble [([Edge], Point v n -> Point v n -> v n)]
fs Map PID (Particle v n)
ps)
          = [([Edge], Point v n -> Point v n -> v n)]
-> Map PID (Particle v n) -> Ensemble v n
forall (v :: * -> *) n.
[([Edge], Point v n -> Point v n -> v n)]
-> Map PID (Particle v n) -> Ensemble v n
Ensemble [([Edge], Point v n -> Point v n -> v n)]
fs
            ((Unwrapped (Endo (Map PID (Particle v n)))
 -> Endo (Map PID (Particle v n)))
-> ((Unwrapped (Endo (Map PID (Particle v n)))
     -> Endo (Map PID (Particle v n)))
    -> [Map PID (Particle v n) -> Map PID (Particle v n)]
    -> Endo (Map PID (Particle v n)))
-> [Map PID (Particle v n) -> Map PID (Particle v n)]
-> Unwrapped (Endo (Map PID (Particle v n)))
forall (f :: * -> *) s t.
(Functor f, Rewrapping s t) =>
(Unwrapped s -> s)
-> ((Unwrapped t -> t) -> f s) -> f (Unwrapped s)
ala Unwrapped (Endo (Map PID (Particle v n)))
-> Endo (Map PID (Particle v n))
(Map PID (Particle v n) -> Map PID (Particle v n))
-> Endo (Map PID (Particle v n))
forall a. (a -> a) -> Endo a
Endo (Unwrapped (Endo (Map PID (Particle v n)))
 -> Endo (Map PID (Particle v n)))
-> [Unwrapped (Endo (Map PID (Particle v n)))]
-> Endo (Map PID (Particle v n))
(Unwrapped (Endo (Map PID (Particle v n)))
 -> Endo (Map PID (Particle v n)))
-> [Map PID (Particle v n) -> Map PID (Particle v n)]
-> Endo (Map PID (Particle v n))
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap ((([Edge], Point v n -> Point v n -> v n)
 -> [Map PID (Particle v n) -> Map PID (Particle v n)])
-> [([Edge], Point v n -> Point v n -> v n)]
-> [Map PID (Particle v n) -> Map PID (Particle v n)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\([Edge]
es, Point v n -> Point v n -> v n
f) -> ((Edge -> Map PID (Particle v n) -> Map PID (Particle v n))
-> [Edge] -> [Map PID (Particle v n) -> Map PID (Particle v n)]
forall a b. (a -> b) -> [a] -> [b]
map ((Point v n -> Point v n -> v n)
-> Edge -> Map PID (Particle v n) -> Map PID (Particle v n)
forall {k} {v :: * -> *} {n}.
(Ord k, Additive v, Num n) =>
(Point v n -> Point v n -> v n)
-> (k, k) -> Map k (Particle v n) -> Map k (Particle v n)
mkForce Point v n -> Point v n -> v n
f) [Edge]
es)) [([Edge], Point v n -> Point v n -> v n)]
fs) Map PID (Particle v n)
ps)
        -- mkForce :: (Point v n -> Point v n -> v n)
        --         -> Edge
        --         -> M.Map Int (Particle v n)
        --         -> M.Map Int (Particle v n)
        mkForce :: (Point v n -> Point v n -> v n)
-> (k, k) -> Map k (Particle v n) -> Map k (Particle v n)
mkForce Point v n -> Point v n -> v n
f (k
i1, k
i2) Map k (Particle v n)
m
          = case (k -> Map k (Particle v n) -> Maybe (Particle v n)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
i1 Map k (Particle v n)
m, k -> Map k (Particle v n) -> Maybe (Particle v n)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
i2 Map k (Particle v n)
m) of
              (Just Particle v n
p1, Just Particle v n
p2) ->
                ( (Particle v n -> Particle v n)
-> k -> Map k (Particle v n) -> Map k (Particle v n)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust ((v n -> Identity (v n)) -> Particle v n -> Identity (Particle v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(v n -> f (v n)) -> Particle v n -> f (Particle v n)
force ((v n -> Identity (v n))
 -> Particle v n -> Identity (Particle v n))
-> (v n -> v n) -> Particle v n -> Particle v n
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (v n -> v n -> v n
forall a. Num a => v a -> v a -> v a
forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a
^+^ Point v n -> Point v n -> v n
f (Particle v n
p1Particle v n
-> Getting (Point v n) (Particle v n) (Point v n) -> Point v n
forall s a. s -> Getting a s a -> a
^.Getting (Point v n) (Particle v n) (Point v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(Point v n -> f (Point v n)) -> Particle v n -> f (Particle v n)
pos) (Particle v n
p2Particle v n
-> Getting (Point v n) (Particle v n) (Point v n) -> Point v n
forall s a. s -> Getting a s a -> a
^.Getting (Point v n) (Particle v n) (Point v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(Point v n -> f (Point v n)) -> Particle v n -> f (Particle v n)
pos))) k
i1
                (Map k (Particle v n) -> Map k (Particle v n))
-> (Map k (Particle v n) -> Map k (Particle v n))
-> Map k (Particle v n)
-> Map k (Particle v n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Particle v n -> Particle v n)
-> k -> Map k (Particle v n) -> Map k (Particle v n)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust ((v n -> Identity (v n)) -> Particle v n -> Identity (Particle v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(v n -> f (v n)) -> Particle v n -> f (Particle v n)
force ((v n -> Identity (v n))
 -> Particle v n -> Identity (Particle v n))
-> (v n -> v n) -> Particle v n -> Particle v n
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (v n -> v n -> v n
forall a. Num a => v a -> v a -> v a
forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a
^-^ Point v n -> Point v n -> v n
f (Particle v n
p1Particle v n
-> Getting (Point v n) (Particle v n) (Point v n) -> Point v n
forall s a. s -> Getting a s a -> a
^.Getting (Point v n) (Particle v n) (Point v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(Point v n -> f (Point v n)) -> Particle v n -> f (Particle v n)
pos) (Particle v n
p2Particle v n
-> Getting (Point v n) (Particle v n) (Point v n) -> Point v n
forall s a. s -> Getting a s a -> a
^.Getting (Point v n) (Particle v n) (Point v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(Point v n -> f (Point v n)) -> Particle v n -> f (Particle v n)
pos))) k
i2)
                Map k (Particle v n)
m
              (Maybe (Particle v n), Maybe (Particle v n))
_                  -> Map k (Particle v n)
m

-- | Compute the total kinetic energy of an ensemble.
kineticEnergy :: (Metric v, Num n) => Ensemble v n -> n
kineticEnergy :: forall (v :: * -> *) n. (Metric v, Num n) => Ensemble v n -> n
kineticEnergy = (Unwrapped (Sum n) -> Sum n)
-> ((Unwrapped (Sum n) -> Sum n) -> Map PID n -> Sum n)
-> Map PID n
-> Unwrapped (Sum n)
forall (f :: * -> *) s t.
(Functor f, Rewrapping s t) =>
(Unwrapped s -> s)
-> ((Unwrapped t -> t) -> f s) -> f (Unwrapped s)
ala n -> Sum n
Unwrapped (Sum n) -> Sum n
forall a. a -> Sum a
Sum (Unwrapped (Sum n) -> Sum n) -> Map PID n -> Sum n
(Unwrapped (Sum n) -> Sum n)
-> Map PID (Unwrapped (Sum n)) -> Sum n
forall m a. Monoid m => (a -> m) -> Map PID a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap (Map PID n -> n)
-> (Ensemble v n -> Map PID n) -> Ensemble v n -> n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Particle v n -> n) -> Map PID (Particle v n) -> Map PID n
forall a b. (a -> b) -> Map PID a -> Map PID b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (v n -> n
forall a. Num a => v a -> a
forall (f :: * -> *) a. (Metric f, Num a) => f a -> a
quadrance (v n -> n) -> (Particle v n -> v n) -> Particle v n -> n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (v n) (Particle v n) (v n) -> Particle v n -> v n
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting (v n) (Particle v n) (v n)
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(v n -> f (v n)) -> Particle v n -> f (Particle v n)
vel) (Map PID (Particle v n) -> Map PID n)
-> (Ensemble v n -> Map PID (Particle v n))
-> Ensemble v n
-> Map PID n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting
  (Map PID (Particle v n)) (Ensemble v n) (Map PID (Particle v n))
-> Ensemble v n -> Map PID (Particle v n)
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting
  (Map PID (Particle v n)) (Ensemble v n) (Map PID (Particle v n))
forall (v :: * -> *) n (f :: * -> *).
Functor f =>
(Map PID (Particle v n) -> f (Map PID (Particle v n)))
-> Ensemble v n -> f (Ensemble v n)
particles

------------------------------------------------------------
--  Simulation
------------------------------------------------------------

-- | Options for customizing a simulation.
data ForceLayoutOpts n =
  FLOpts
  { forall n. ForceLayoutOpts n -> n
_damping     :: n         -- ^ Damping factor to be
                              --   applied at each step.
                              --   Should be between 0 and 1.
                              --   The default is 0.8.
  , forall n. ForceLayoutOpts n -> Maybe n
_energyLimit :: Maybe n   -- ^ Kinetic energy below which
                              --   simulation should stop.
                              --   If @Nothing@, pay no
                              --   attention to kinetic
                              --   energy.  The default is
                              --   @Just 0.001@.
  , forall n. ForceLayoutOpts n -> Maybe PID
_stepLimit   :: Maybe Int -- ^ Maximum number of
                              --   simulation steps.  If
                              --   @Nothing@, pay no
                              --   attention to the number of
                              --   steps.  The default is
                              --   @Just 1000@.
  }

makeLenses ''ForceLayoutOpts

instance Fractional n => Default (ForceLayoutOpts n) where
  def :: ForceLayoutOpts n
def = FLOpts
        { _damping :: n
_damping     = n
0.8
        , _energyLimit :: Maybe n
_energyLimit = n -> Maybe n
forall a. a -> Maybe a
Just n
0.001
        , _stepLimit :: Maybe PID
_stepLimit   = PID -> Maybe PID
forall a. a -> Maybe a
Just PID
1000
        }

-- | Simulate a starting ensemble according to the given options,
--   producing a list of all the intermediate ensembles.  Useful for,
--   /e.g./, making an animation.  Note that the resulting list could
--   be infinite, if a 'stepLimit' is not specified and either the
--   kinetic energy never falls below the specified threshold, or no
--   energy threshold is specified.
simulate :: (Metric v, Num n, Ord n)
         => ForceLayoutOpts n -> Ensemble v n -> [Ensemble v n]
simulate :: forall (v :: * -> *) n.
(Metric v, Num n, Ord n) =>
ForceLayoutOpts n -> Ensemble v n -> [Ensemble v n]
simulate ForceLayoutOpts n
opts Ensemble v n
e
  = (Ensemble v n
eEnsemble v n -> [Ensemble v n] -> [Ensemble v n]
forall a. a -> [a] -> [a]
:)
  ([Ensemble v n] -> [Ensemble v n])
-> (Ensemble v n -> [Ensemble v n])
-> Ensemble v n
-> [Ensemble v n]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Ensemble v n -> Bool) -> [Ensemble v n] -> [Ensemble v n]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((n -> Bool) -> (n -> n -> Bool) -> Maybe n -> n -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Bool -> n -> Bool
forall a b. a -> b -> a
const Bool
True) n -> n -> Bool
forall a. Ord a => a -> a -> Bool
(<) (ForceLayoutOpts n
opts ForceLayoutOpts n
-> Getting (Maybe n) (ForceLayoutOpts n) (Maybe n) -> Maybe n
forall s a. s -> Getting a s a -> a
^. Getting (Maybe n) (ForceLayoutOpts n) (Maybe n)
forall n (f :: * -> *).
Functor f =>
(Maybe n -> f (Maybe n))
-> ForceLayoutOpts n -> f (ForceLayoutOpts n)
energyLimit) (n -> Bool) -> (Ensemble v n -> n) -> Ensemble v n -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ensemble v n -> n
forall (v :: * -> *) n. (Metric v, Num n) => Ensemble v n -> n
kineticEnergy)
  ([Ensemble v n] -> [Ensemble v n])
-> (Ensemble v n -> [Ensemble v n])
-> Ensemble v n
-> [Ensemble v n]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Ensemble v n] -> [Ensemble v n])
-> (PID -> [Ensemble v n] -> [Ensemble v n])
-> Maybe PID
-> [Ensemble v n]
-> [Ensemble v n]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [Ensemble v n] -> [Ensemble v n]
forall a. a -> a
id PID -> [Ensemble v n] -> [Ensemble v n]
forall a. PID -> [a] -> [a]
take (ForceLayoutOpts n
opts ForceLayoutOpts n
-> Getting (Maybe PID) (ForceLayoutOpts n) (Maybe PID) -> Maybe PID
forall s a. s -> Getting a s a -> a
^. Getting (Maybe PID) (ForceLayoutOpts n) (Maybe PID)
forall n (f :: * -> *).
Functor f =>
(Maybe PID -> f (Maybe PID))
-> ForceLayoutOpts n -> f (ForceLayoutOpts n)
stepLimit)
  ([Ensemble v n] -> [Ensemble v n])
-> (Ensemble v n -> [Ensemble v n])
-> Ensemble v n
-> [Ensemble v n]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PID -> [Ensemble v n] -> [Ensemble v n]
forall a. PID -> [a] -> [a]
drop PID
1
  ([Ensemble v n] -> [Ensemble v n])
-> (Ensemble v n -> [Ensemble v n])
-> Ensemble v n
-> [Ensemble v n]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Ensemble v n -> Ensemble v n) -> Ensemble v n -> [Ensemble v n]
forall a. (a -> a) -> a -> [a]
iterate (n -> Ensemble v n -> Ensemble v n
forall (v :: * -> *) n.
(Additive v, Num n) =>
n -> Ensemble v n -> Ensemble v n
ensembleStep (ForceLayoutOpts n
opts ForceLayoutOpts n -> Getting n (ForceLayoutOpts n) n -> n
forall s a. s -> Getting a s a -> a
^. Getting n (ForceLayoutOpts n) n
forall n (f :: * -> *).
Functor f =>
(n -> f n) -> ForceLayoutOpts n -> f (ForceLayoutOpts n)
damping))
  (Ensemble v n -> [Ensemble v n]) -> Ensemble v n -> [Ensemble v n]
forall a b. (a -> b) -> a -> b
$ Ensemble v n
e

-- | Run a simluation from a starting ensemble, yielding either the
--   first ensemble to have kinetic energy below the 'energyLimit' (if
--   given), or the ensemble that results after a number of steps
--   equal to the 'stepLimit' (if given), whichever comes first.
--   Otherwise @forceLayout@ will not terminate.
forceLayout :: (Metric v, Num n, Ord n)
            => ForceLayoutOpts n -> Ensemble v n -> Ensemble v n
forceLayout :: forall (v :: * -> *) n.
(Metric v, Num n, Ord n) =>
ForceLayoutOpts n -> Ensemble v n -> Ensemble v n
forceLayout ForceLayoutOpts n
opts = [Ensemble v n] -> Ensemble v n
forall a. HasCallStack => [a] -> a
last ([Ensemble v n] -> Ensemble v n)
-> (Ensemble v n -> [Ensemble v n]) -> Ensemble v n -> Ensemble v n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForceLayoutOpts n -> Ensemble v n -> [Ensemble v n]
forall (v :: * -> *) n.
(Metric v, Num n, Ord n) =>
ForceLayoutOpts n -> Ensemble v n -> [Ensemble v n]
simulate ForceLayoutOpts n
opts

------------------------------------------------------------
--  Standard forces
------------------------------------------------------------

-- | @distForce f p1 p2@ computes the force between two points as a
--   multiple of the unit vector in the direction from @p1@ to @p2@,
--   given a function @f@ which computes the force's magnitude as a
--   function of the distance between the points.
distForce :: (Metric v, Floating n) => (n -> n) -> Point v n -> Point v n -> v n
distForce :: forall (v :: * -> *) n.
(Metric v, Floating n) =>
(n -> n) -> Point v n -> Point v n -> v n
distForce n -> n
f Point v n
p1 Point v n
p2 = n -> v n -> v n
forall {f :: * -> *} {a}. (Metric f, Floating a) => a -> f a -> f a
withLength (n -> n
f (Point v n -> Point v n -> n
forall a. Floating a => Point v a -> Point v a -> a
forall (f :: * -> *) a. (Metric f, Floating a) => f a -> f a -> a
distance Point v n
p1 Point v n
p2)) (Point v n
p2 Point v n -> Point v n -> Diff (Point v) n
forall a. Num a => Point v a -> Point v a -> Diff (Point v) a
forall (p :: * -> *) a. (Affine p, Num a) => p a -> p a -> Diff p a
.-. Point v n
p1)
  where withLength :: a -> f a -> f a
withLength a
s f a
v = a
s a -> f a -> f a
forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ f a -> f a
forall a. Floating a => f a -> f a
forall (f :: * -> *) a. (Metric f, Floating a) => f a -> f a
signorm f a
v

-- | @hookeForce k l p1 p2@ computes the force on @p1@, assuming that
--   @p1@ and @p2@ are connected by a spring with equilibrium length @l@
--   and spring constant @k@.
hookeForce :: (Metric v, Floating n) => n -> n -> Point v n -> Point v n -> v n
hookeForce :: forall (v :: * -> *) n.
(Metric v, Floating n) =>
n -> n -> Point v n -> Point v n -> v n
hookeForce n
k n
l = (n -> n) -> Point v n -> Point v n -> v n
forall (v :: * -> *) n.
(Metric v, Floating n) =>
(n -> n) -> Point v n -> Point v n -> v n
distForce (\n
d -> n
k n -> n -> n
forall a. Num a => a -> a -> a
* (n
d n -> n -> n
forall a. Num a => a -> a -> a
- n
l))

-- | @coulombForce k@ computes the electrostatic repulsive force
--   between two charged particles, with constant of proportionality
--   @k@.
coulombForce :: (Metric v, Floating n) => n -> Point v n -> Point v n -> v n
coulombForce :: forall (v :: * -> *) n.
(Metric v, Floating n) =>
n -> Point v n -> Point v n -> v n
coulombForce n
k = (n -> n) -> Point v n -> Point v n -> v n
forall (v :: * -> *) n.
(Metric v, Floating n) =>
(n -> n) -> Point v n -> Point v n -> v n
distForce (\n
d -> -n
kn -> n -> n
forall a. Fractional a => a -> a -> a
/(n
dn -> n -> n
forall a. Num a => a -> a -> a
*n
d))