uvector-0.1: examples/barhesHut/BarnesHutSeq.hs
{-# GHC_OPTIONS -fglasgow-exts #-}
module BarnesHutSeq
where
import Data.Array.Parallel.Unlifted
import BarnesHutGen
-- Phase 1: building the tree
--
{-
-- Split massPoints according to their locations in the quadrants
--
splitPoints:: BoundingBox -> UArr MassPoint -> SUArr MassPoint
splitPoints (ll@(llx :*: lly) :*: ru@(rux :*: ruy)) particles
| noOfPoints == 0 = singletonSU particles
| otherwise = singletonSU lls +:+^ singletonSU lus +:+^ singletonSU rus +:+^ singletonSU rls
where
noOfPoints = lengthU particles
lls = filterU (inBox (ll :*: mid)) particles
lus = filterU (inBox ((llx :*: midy) :*: (midx :*: ruy ))) particles
rus = filterU (inBox (mid :*: ru )) particles
rls = filterU (inBox ((midx :*: lly) :*: (rux :*: midy))) particles
mid@(midx :*: midy) = ((llx + rux)/2.0) :*: ((lly + ruy)/2.0)
-}
splitPointsL:: UArr BoundingBox -> SUArr MassPoint -> BHTree
splitPointsL bboxes particless
| lengthSU multiparticles == 0 = [(centroids, toU [])]
| otherwise = (centroids, lengthsSU multiparticles) :
(splitPointsL newBoxes multiparticles)
where
-- calculate centroid of each segment
centroids =
calcCentroids $ segmentArrU nonEmptySegd $ flattenSU particless
-- remove empty segments
multiPointFlags = mapU ((>1)) $ lengthsSU particless
multiparticles = (splitPointsL' llbb lubb rubb rlbb) $
packCU multiPointFlags particless
bboxes' = packU bboxes multiPointFlags
nonEmptySegd = filterU ((>0)) $ lengthsSU particless
-- split each box in four sub-boxes
newBoxes = merge4 llbb lubb rubb rlbb
llbb = mapU makells bboxes'
lubb = mapU makelus bboxes'
rubb = mapU makerus bboxes'
rlbb = mapU makerls bboxes'
makells (ll@(llx :*: lly) :*: ru@(rux :*: ruy)) =
ll :*: (((llx + rux)/2.0) :*: (((lly + ruy)/2.0)))
makelus (ll@(llx :*: lly) :*: ru@(rux :*: ruy)) =
(llx :*: ((lly + ruy)/2.0)) :*: (((llx + rux)/2.0) :*: ruy )
makerus (ll@(llx :*: lly) :*: ru@(rux :*: ruy)) =
(((llx + rux)/2.0) :*: ((lly + ruy)/2.0)) :*: ru
makerls (ll@(llx :*: lly) :*: ru@(rux :*: ruy)) =
((((llx + rux)/2.0) :*: lly) :*: (rux :*: ((lly + ruy)/2.0)))
splitPointsL':: UArr BoundingBox ->
UArr BoundingBox ->
UArr BoundingBox ->
UArr BoundingBox ->
SUArr MassPoint ->
SUArr MassPoint
splitPointsL' llbb lubb rubb rlbb particless
| particlessLen == 0 = particless
| otherwise = orderedPoints
where
-- each segment split into four subsegments with particles located in
-- the four quadrants
orderedPoints =
segmentArrU newLengths $
flattenSU $ llsPs ^+:+^ lusPs ^+:+^ rusPs ^+:+^ rlsPs
particlessLen = lengthSU particless
pssLens = lengthsSU particless
lls = replicateSU pssLens llbb
lus = replicateSU pssLens lubb
rus = replicateSU pssLens rubb
rls = replicateSU pssLens rlbb
llsPs = mapSU sndS $ filterSU (uncurryS inBox)
(zipSU (replicateSU pssLens llbb) particless)
lusPs = mapSU sndS $ filterSU (uncurryS inBox)
(zipSU (replicateSU pssLens lubb) particless)
rusPs = mapSU sndS $ filterSU (uncurryS inBox)
(zipSU (replicateSU pssLens rubb) particless)
rlsPs = mapSU sndS $ filterSU (uncurryS inBox)
(zipSU (replicateSU pssLens rlbb) particless)
newLengths =
merge4 (lengthsSU llsPs) (lengthsSU lusPs)
(lengthsSU rusPs) (lengthsSU rlsPs)
-- Calculate centroid of each subarray
--
calcCentroids:: SUArr MassPoint -> UArr MassPoint
calcCentroids orderedPoints = centroids
where
ms = foldSU (+) 0.0 $ sndSU orderedPoints
centroids = zipWithU div' ms $
foldSU pairP (0.0 :*: 0.0) $
zipWithSU multCoor orderedPoints
(replicateSU (lengthsSU orderedPoints) ms)
div' m (x :*: y) = ((x/m :*: y/m) :*: m)
multCoor ((x :*: y) :*: _) m = (m * x :*: m * y)
pairP (x1 :*: y1) (x2 :*: y2) = ((x1+x2) :*: (y1 + y2))
-- phase 2:
-- calculating the velocities
calcAccel:: BHTree -> UArr MassPoint -> UArr (Double :*: Double)
calcAccel [] particles
| lengthU particles == 0 = emptyU
| otherwise = error $ "calcVelocity: reached empty tree" ++ (show particles)
calcAccel ((centroids, segd) :trees) particles = closeAccel
where
closeAccel = splitApplyU particlesClose
((calcAccel trees) . sndU )
calcFarAccel
(zipU
(flattenSU $ replicateCU (lengthU particles) centroids)
(flattenSU $ replicateSU segd particles))
particlesClose (((x1 :*: y1):*: _) :*: ((x2 :*: y2) :*: _)) =
(x1-x2)^2 + (y1-y2)^2 < eClose
calcFarAccel:: UArr (MassPoint :*: MassPoint) -> UArr Accel
calcFarAccel = mapU accel
--
--
accel:: MassPoint :*: MassPoint -> Accel
accel (((x1:*: y1) :*: m) :*:
((x2:*: y2) :*: _)) | r < epsilon = (0.0 :*: 0.0)
| otherwise = (aabs * dx / r :*: aabs * dy / r)
where
rsqr = (dx * dx) + (dy * dy)
r = sqrt rsqr
dx = x1 - x2
dy = y1 - y2
aabs = m / rsqr
-- assumes all arr have the same length
-- result [a11, a21, a31, a41, a12, a22....]
merge4:: UA a =>
UArr a ->UArr a ->UArr a ->UArr a ->UArr a
merge4 a1 a2 a3 a4 =
combineU flags3 (combineU flags2 (combineU flags1 a1 a2) a3) a4
where
flags1 = mapU even $ enumFromToU 0 (2 * len-1)
flags2 = mapU (\x -> mod x 3 /= 2) $ enumFromToU 0 (3 * len-1)
flags3 = mapU (\x -> mod x 4 /= 3) $ enumFromToU 0 (4 * len-1)
len = lengthU a1
-- checks if particle is in box (excluding left and lower border)
inBox:: BoundingBox -> MassPoint -> Bool
inBox ((ll@(llx :*: lly) :*: ru@(rux :*: ruy))) ((px :*: py) :*: _) =
(px > llx) && (px <= rux) && (py > lly) && (py <= ruy)
splitApplyU:: (UA e, UA e') => (e -> Bool) -> (UArr e -> UArr e') -> (UArr e -> UArr e') -> UArr e -> UArr e'
splitApplyU p f1 f2 xsArr = combineU (mapU p xsArr) res1 res2
where
res1 = f1 $ filterU p xsArr
res2 = f2 $ filterU (not . p) xsArr
splitApplySU:: (UA e, UA e') => UArr Bool -> (SUArr e -> SUArr e') -> (SUArr e -> SUArr e') -> SUArr e -> SUArr e'
{-# INLINE splitApplySU #-}
splitApplySU flags f1 f2 xssArr = combineCU flags res1 res2
where
res1 = f1 $ packCU flags xssArr
res2 = f2 $ packCU (mapU not flags) xssArr