SpringModel.hs

module HGraph.Undirected.Layout.SpringModel
        (setup, step, positions)
where

import HGraph.Undirected
import qualified Data.Array.Unboxed as U
import qualified Data.Map as M
import Data.Foldable
import Linear

data SpringModel g a = 
  SpringModel
  { smGraph :: g Int
  , smPoints :: U.Array Int (V2 Double)
  , smVelocities :: U.Array Int (V2 Double)
  , smSpringStrength    :: Double
  , smRepulsionStrength :: Double
  , smFriction :: Double
  , smCenterAttractionStrength :: Double
  , smIndexToLabel :: M.Map Int a
  }

setup centerAttraction springStrength repulsiveStrength friction gr =
  let (gi, itol) = linearizeVertices gr
      n = numVertices gi - 1
  in SpringModel
  { smGraph = gi
  , smPoints = U.array (0, n) $ zip [0..n] (repeat 0)
  , smVelocities = U.array (0, n) $ zip [0..n] (repeat 0)
  , smSpringStrength = springStrength
  , smRepulsionStrength = repulsiveStrength
  , smFriction = friction
  , smCenterAttractionStrength = centerAttraction
  , smIndexToLabel = M.fromList itol
  }

step dt layout = layout{smPoints = points', smVelocities = velocities'}
  where
    points = smPoints layout
    velocities = smVelocities layout
    points' = U.array (U.bounds points)
      [ (i, (points U.! i) + (velocities U.! i))
      | i <- U.indices points 
      ]
    velocities' = U.array (U.bounds points)
                          [ (i, ((smFriction layout) ** dt) *^ (velocities U.! i) + (accelerate layout i))
                          | i <- U.indices $ points]

accelerate layout v = springV ^+^ centerV ^+^ pushV
  where
    positions = smPoints layout
    pv = positions U.! v
    springV = sum
      [ du ^* (smSpringStrength layout)
      | u <- neighbors (smGraph layout) v
      , let pu = positions U.! u
      , let du = pu ^-^ pv
      ]
    centerV = (negated pv) ^* (smCenterAttractionStrength layout)
    pushV = sum
      [ if d2 < 0.0001 then
           (100 * smRepulsionStrength layout) *^ (angle $ 2*pi * ((fromIntegral v) / (fromIntegral $ U.rangeSize $ U.bounds positions)))
        else
           ((du ^/ (sqrt d2)) ^* (smRepulsionStrength layout)) ^/ d2
      | u <- vertices $ smGraph layout
      , let pu = positions U.! u
      , let du = pv ^-^ pu
      , let d2 = foldl' (\s x -> s + x*x) 0 du
      , u /= v
      ]

positions layout =
  [ (smIndexToLabel layout M.! i, (smPoints layout) U.! i)
  | i <- U.indices $ smPoints layout
  ]