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 { SpringModel g a -> g Int smGraph :: g Int , SpringModel g a -> Array Int (V2 Double) smPoints :: U.Array Int (V2 Double) , SpringModel g a -> Array Int (V2 Double) smVelocities :: U.Array Int (V2 Double) , SpringModel g a -> Double smSpringStrength :: Double , SpringModel g a -> Double smRepulsionStrength :: Double , SpringModel g a -> Double smFriction :: Double , SpringModel g a -> Double smCenterAttractionStrength :: Double , SpringModel g a -> Map Int a smIndexToLabel :: M.Map Int a } setup :: Double -> Double -> Double -> Double -> g a -> SpringModel g a setup Double centerAttraction Double springStrength Double repulsiveStrength Double friction g a gr = let (g Int gi, [(Int, a)] itol) = g a -> (g Int, [(Int, a)]) forall (t :: * -> *) a. UndirectedGraph t => t a -> (t Int, [(Int, a)]) linearizeVertices g a gr n :: Int n = g Int -> Int forall (t :: * -> *) b a. (UndirectedGraph t, Integral b) => t a -> b numVertices g Int gi Int -> Int -> Int forall a. Num a => a -> a -> a - Int 1 in SpringModel :: forall (g :: * -> *) a. g Int -> Array Int (V2 Double) -> Array Int (V2 Double) -> Double -> Double -> Double -> Double -> Map Int a -> SpringModel g a SpringModel { smGraph :: g Int smGraph = g Int gi , smPoints :: Array Int (V2 Double) smPoints = (Int, Int) -> [(Int, V2 Double)] -> Array Int (V2 Double) forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => (i, i) -> [(i, e)] -> a i e U.array (Int 0, Int n) ([(Int, V2 Double)] -> Array Int (V2 Double)) -> [(Int, V2 Double)] -> Array Int (V2 Double) forall a b. (a -> b) -> a -> b $ [Int] -> [V2 Double] -> [(Int, V2 Double)] forall a b. [a] -> [b] -> [(a, b)] zip [Int 0..Int n] (V2 Double -> [V2 Double] forall a. a -> [a] repeat V2 Double 0) , smVelocities :: Array Int (V2 Double) smVelocities = (Int, Int) -> [(Int, V2 Double)] -> Array Int (V2 Double) forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => (i, i) -> [(i, e)] -> a i e U.array (Int 0, Int n) ([(Int, V2 Double)] -> Array Int (V2 Double)) -> [(Int, V2 Double)] -> Array Int (V2 Double) forall a b. (a -> b) -> a -> b $ [Int] -> [V2 Double] -> [(Int, V2 Double)] forall a b. [a] -> [b] -> [(a, b)] zip [Int 0..Int n] (V2 Double -> [V2 Double] forall a. a -> [a] repeat V2 Double 0) , smSpringStrength :: Double smSpringStrength = Double springStrength , smRepulsionStrength :: Double smRepulsionStrength = Double repulsiveStrength , smFriction :: Double smFriction = Double friction , smCenterAttractionStrength :: Double smCenterAttractionStrength = Double centerAttraction , smIndexToLabel :: Map Int a smIndexToLabel = [(Int, a)] -> Map Int a forall k a. Ord k => [(k, a)] -> Map k a M.fromList [(Int, a)] itol } step :: Double -> SpringModel g a -> SpringModel g a step Double dt SpringModel g a layout = SpringModel g a layout{smPoints :: Array Int (V2 Double) smPoints = Array Int (V2 Double) points', smVelocities :: Array Int (V2 Double) smVelocities = Array Int (V2 Double) velocities'} where points :: Array Int (V2 Double) points = SpringModel g a -> Array Int (V2 Double) forall (g :: * -> *) a. SpringModel g a -> Array Int (V2 Double) smPoints SpringModel g a layout velocities :: Array Int (V2 Double) velocities = SpringModel g a -> Array Int (V2 Double) forall (g :: * -> *) a. SpringModel g a -> Array Int (V2 Double) smVelocities SpringModel g a layout points' :: Array Int (V2 Double) points' = (Int, Int) -> [(Int, V2 Double)] -> Array Int (V2 Double) forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => (i, i) -> [(i, e)] -> a i e U.array (Array Int (V2 Double) -> (Int, Int) forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> (i, i) U.bounds Array Int (V2 Double) points) [ (Int i, (Array Int (V2 Double) points Array Int (V2 Double) -> Int -> V2 Double forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> i -> e U.! Int i) V2 Double -> V2 Double -> V2 Double forall a. Num a => a -> a -> a + (Array Int (V2 Double) velocities Array Int (V2 Double) -> Int -> V2 Double forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> i -> e U.! Int i)) | Int i <- Array Int (V2 Double) -> [Int] forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [i] U.indices Array Int (V2 Double) points ] velocities' :: Array Int (V2 Double) velocities' = (Int, Int) -> [(Int, V2 Double)] -> Array Int (V2 Double) forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => (i, i) -> [(i, e)] -> a i e U.array (Array Int (V2 Double) -> (Int, Int) forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> (i, i) U.bounds Array Int (V2 Double) points) [ (Int i, ((SpringModel g a -> Double forall (g :: * -> *) a. SpringModel g a -> Double smFriction SpringModel g a layout) Double -> Double -> Double forall a. Floating a => a -> a -> a ** Double dt) Double -> V2 Double -> V2 Double forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a *^ (Array Int (V2 Double) velocities Array Int (V2 Double) -> Int -> V2 Double forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> i -> e U.! Int i) V2 Double -> V2 Double -> V2 Double forall a. Num a => a -> a -> a + (SpringModel g a -> Int -> V2 Double forall (g :: * -> *) a. Adjacency g => SpringModel g a -> Int -> V2 Double accelerate SpringModel g a layout Int i)) | Int i <- Array Int (V2 Double) -> [Int] forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [i] U.indices (Array Int (V2 Double) -> [Int]) -> Array Int (V2 Double) -> [Int] forall a b. (a -> b) -> a -> b $ Array Int (V2 Double) points] accelerate :: SpringModel g a -> Int -> V2 Double accelerate SpringModel g a layout Int v = V2 Double springV V2 Double -> V2 Double -> V2 Double forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a ^+^ V2 Double centerV V2 Double -> V2 Double -> V2 Double forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a ^+^ V2 Double pushV where positions :: Array Int (V2 Double) positions = SpringModel g a -> Array Int (V2 Double) forall (g :: * -> *) a. SpringModel g a -> Array Int (V2 Double) smPoints SpringModel g a layout pv :: V2 Double pv = Array Int (V2 Double) positions Array Int (V2 Double) -> Int -> V2 Double forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> i -> e U.! Int v springV :: V2 Double springV = [V2 Double] -> V2 Double forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a sum [ V2 Double du V2 Double -> Double -> V2 Double forall (f :: * -> *) a. (Functor f, Num a) => f a -> a -> f a ^* (SpringModel g a -> Double forall (g :: * -> *) a. SpringModel g a -> Double smSpringStrength SpringModel g a layout) | Int u <- g Int -> Int -> [Int] forall (t :: * -> *) a. Adjacency t => t a -> a -> [a] neighbors (SpringModel g a -> g Int forall (g :: * -> *) a. SpringModel g a -> g Int smGraph SpringModel g a layout) Int v , let pu :: V2 Double pu = Array Int (V2 Double) positions Array Int (V2 Double) -> Int -> V2 Double forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> i -> e U.! Int u , let du :: V2 Double du = V2 Double pu V2 Double -> V2 Double -> V2 Double forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a ^-^ V2 Double pv ] centerV :: V2 Double centerV = (V2 Double -> V2 Double forall (f :: * -> *) a. (Functor f, Num a) => f a -> f a negated V2 Double pv) V2 Double -> Double -> V2 Double forall (f :: * -> *) a. (Functor f, Num a) => f a -> a -> f a ^* (SpringModel g a -> Double forall (g :: * -> *) a. SpringModel g a -> Double smCenterAttractionStrength SpringModel g a layout) pushV :: V2 Double pushV = [V2 Double] -> V2 Double forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a sum [ if Double d2 Double -> Double -> Bool forall a. Ord a => a -> a -> Bool < Double 0.0001 then (Double 100 Double -> Double -> Double forall a. Num a => a -> a -> a * SpringModel g a -> Double forall (g :: * -> *) a. SpringModel g a -> Double smRepulsionStrength SpringModel g a layout) Double -> V2 Double -> V2 Double forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a *^ (Double -> V2 Double forall a. Floating a => a -> V2 a angle (Double -> V2 Double) -> Double -> V2 Double forall a b. (a -> b) -> a -> b $ Double 2Double -> Double -> Double forall a. Num a => a -> a -> a *Double forall a. Floating a => a pi Double -> Double -> Double forall a. Num a => a -> a -> a * ((Int -> Double forall a b. (Integral a, Num b) => a -> b fromIntegral Int v) Double -> Double -> Double forall a. Fractional a => a -> a -> a / (Int -> Double forall a b. (Integral a, Num b) => a -> b fromIntegral (Int -> Double) -> Int -> Double forall a b. (a -> b) -> a -> b $ (Int, Int) -> Int forall a. Ix a => (a, a) -> Int U.rangeSize ((Int, Int) -> Int) -> (Int, Int) -> Int forall a b. (a -> b) -> a -> b $ Array Int (V2 Double) -> (Int, Int) forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> (i, i) U.bounds Array Int (V2 Double) positions))) else ((V2 Double du V2 Double -> Double -> V2 Double forall (f :: * -> *) a. (Functor f, Fractional a) => f a -> a -> f a ^/ (Double -> Double forall a. Floating a => a -> a sqrt Double d2)) V2 Double -> Double -> V2 Double forall (f :: * -> *) a. (Functor f, Num a) => f a -> a -> f a ^* (SpringModel g a -> Double forall (g :: * -> *) a. SpringModel g a -> Double smRepulsionStrength SpringModel g a layout)) V2 Double -> Double -> V2 Double forall (f :: * -> *) a. (Functor f, Fractional a) => f a -> a -> f a ^/ Double d2 | Int u <- g Int -> [Int] forall (t :: * -> *) a. UndirectedGraph t => t a -> [a] vertices (g Int -> [Int]) -> g Int -> [Int] forall a b. (a -> b) -> a -> b $ SpringModel g a -> g Int forall (g :: * -> *) a. SpringModel g a -> g Int smGraph SpringModel g a layout , let pu :: V2 Double pu = Array Int (V2 Double) positions Array Int (V2 Double) -> Int -> V2 Double forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> i -> e U.! Int u , let du :: V2 Double du = V2 Double pv V2 Double -> V2 Double -> V2 Double forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a ^-^ V2 Double pu , let d2 :: Double d2 = (Double -> Double -> Double) -> Double -> V2 Double -> Double forall (t :: * -> *) b a. Foldable t => (b -> a -> b) -> b -> t a -> b foldl' (\Double s Double x -> Double s Double -> Double -> Double forall a. Num a => a -> a -> a + Double xDouble -> Double -> Double forall a. Num a => a -> a -> a *Double x) Double 0 V2 Double du , Int u Int -> Int -> Bool forall a. Eq a => a -> a -> Bool /= Int v ] positions :: SpringModel g a -> [(a, V2 Double)] positions SpringModel g a layout = [ (SpringModel g a -> Map Int a forall (g :: * -> *) a. SpringModel g a -> Map Int a smIndexToLabel SpringModel g a layout Map Int a -> Int -> a forall k a. Ord k => Map k a -> k -> a M.! Int i, (SpringModel g a -> Array Int (V2 Double) forall (g :: * -> *) a. SpringModel g a -> Array Int (V2 Double) smPoints SpringModel g a layout) Array Int (V2 Double) -> Int -> V2 Double forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> i -> e U.! Int i) | Int i <- Array Int (V2 Double) -> [Int] forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [i] U.indices (Array Int (V2 Double) -> [Int]) -> Array Int (V2 Double) -> [Int] forall a b. (a -> b) -> a -> b $ SpringModel g a -> Array Int (V2 Double) forall (g :: * -> *) a. SpringModel g a -> Array Int (V2 Double) smPoints SpringModel g a layout ]