{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} ----------------------------------------------------------------------------- -- | -- 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 Physics.ForceLayout -- > import qualified Data.Map as M -- > import Data.AffineSpace.Point -- > import Data.Default (def) -- > -- > e :: Ensemble (Double, 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) -- > \$ [ (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 (Double, Double) -- > e' = forceLayout (def & damping .~ 0.8 -- > & energyLimit .~ Just 0.001 -- > & stepLimit .~ Nothing -- > ) -- > e -- -- See the diagrams-contrib package -- () 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.AffineSpace import Data.AffineSpace.Point import Data.Default.Class import Data.Foldable (foldMap) import qualified Data.Foldable as F import qualified Data.Map as M import Data.Monoid import Data.VectorSpace hiding (Sum) import Control.Lens ------------------------------------------------------------ -- Particles ------------------------------------------------------------ -- | A particle has a current position, current velocity, and current -- force acting on it. data Particle v = Particle { _pos :: Point v , _vel :: v , _force :: v } deriving (Eq, Show) makeLenses ''Particle -- | Create an initial particle at rest at a particular location. initParticle :: AdditiveGroup v => Point v -> Particle v initParticle p = Particle p zeroV zeroV ------------------------------------------------------------ -- 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 = Ensemble { _forces :: [([Edge], Point v -> Point v -> v)] , _particles :: M.Map PID (Particle v) } makeLenses ''Ensemble ------------------------------------------------------------ -- Simulation internals ------------------------------------------------------------ -- | Simulate one time step for an entire ensemble, with the given -- damping factor. ensembleStep :: VectorSpace v => Scalar v -> Ensemble v -> Ensemble v ensembleStep d = (over particles . M.map) (particleStep d) . recalcForces -- | Simulate one time step for a particle (assuming the force acting -- on it has already been computed), with the given damping factor. particleStep :: VectorSpace v => Scalar v -> Particle v -> Particle v particleStep d = stepPos . stepVel where stepVel p = vel .~ (d *^ (p^.vel ^+^ p^.force)) \$ p stepPos p = pos %~ (.+^ p^.vel) \$ p -- | Recalculate all the forces acting in the next time step of an -- ensemble. recalcForces :: forall v. AdditiveGroup v => Ensemble v -> Ensemble v recalcForces = calcForces . zeroForces where zeroForces = (particles %~) . M.map \$ force .~ zeroV calcForces (Ensemble fs ps) = Ensemble fs (ala Endo foldMap (concatMap (\(es, f) -> (map (mkForce f) es)) fs) ps) mkForce :: (Point v -> Point v -> v) -> Edge -> M.Map Int (Particle v) -> M.Map Int (Particle v) mkForce f (i1, i2) m = case (M.lookup i1 m, M.lookup i2 m) of (Just p1, Just p2) -> ( M.adjust (force %~ (^+^ f (p1^.pos) (p2^.pos))) i1 . M.adjust (force %~ (^-^ f (p1^.pos) (p2^.pos))) i2) m _ -> m -- | Compute the total kinetic energy of an ensemble. kineticEnergy :: (InnerSpace v, Num (Scalar v)) => Ensemble v -> Scalar v kineticEnergy = ala Sum F.foldMap . fmap (magnitudeSq . view vel) . view particles ------------------------------------------------------------ -- Simulation ------------------------------------------------------------ -- | Options for customizing a simulation. data ForceLayoutOpts v = FLOpts { _damping :: Scalar v -- ^ Damping factor to be -- applied at each step. -- Should be between 0 and 1. -- The default is 0.8. , _energyLimit :: Maybe (Scalar v) -- ^ Kinetic energy below which -- simulation should stop. -- If @Nothing@, pay no -- attention to kinetic -- energy. The default is -- @Just 0.001@. , _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 (Scalar v) => Default (ForceLayoutOpts v) where def = FLOpts { _damping = 0.8 , _energyLimit = Just 0.001 , _stepLimit = Just 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 :: (InnerSpace v, Ord (Scalar v), Num (Scalar v)) => ForceLayoutOpts v -> Ensemble v -> [Ensemble v] simulate opts e = (e:) . takeWhile (maybe (const True) (<) (opts ^. energyLimit) . kineticEnergy) . maybe id take (opts ^. stepLimit) . drop 1 . iterate (ensembleStep (opts ^. damping)) \$ 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 :: (InnerSpace v, Ord (Scalar v), Num (Scalar v)) => ForceLayoutOpts v -> Ensemble v -> Ensemble v forceLayout opts = last . simulate 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 :: (InnerSpace v, Floating (Scalar v)) => (Scalar v -> Scalar v) -> Point v -> Point v -> v distForce f p1 p2 = withLength (f (distance p1 p2)) (p2 .-. p1) where withLength s v = s *^ normalized 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 :: (InnerSpace v, Floating (Scalar v)) => Scalar v -> Scalar v -> Point v -> Point v -> v hookeForce k l = distForce (\d -> k * (d - l)) -- | @coulombForce k@ computes the electrostatic repulsive force -- between two charged particles, with constant of proportionality -- @k@. coulombForce :: (InnerSpace v, Floating (Scalar v)) => Scalar v -> Point v -> Point v -> v coulombForce k = distForce (\d -> -k * 1/(d*d))