1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
|
{-# LANGUAGE ParallelArrays #-}
{-# OPTIONS -fvectorise #-}
module Solver
(calcAccelsWithBoxPA)
where
import Data.Array.Parallel
import Data.Array.Parallel.Prelude.Bool
import Data.Array.Parallel.Prelude.Double as D
import qualified Data.Array.Parallel.Prelude.Int as I
import qualified Prelude
data BoundingBox
= Box Double -- lower left X
Double -- lower left Y
Double -- upper right X
Double -- upper right Y
data MassPoint
= MP Double -- pos X
Double -- pos Y
Double -- mass
type Accel
= (Double, Double)
data BHTree
= BHT Double -- size of cell
Double -- centroid X
Double -- centroid Y
Double -- centroid mass
[:BHTree:] -- children
calcAccelsWithBoxPA
:: Double
-> Double -> Double -> Double -> Double
-> PArray (Double, Double, Double)
-> PArray (Double, Double)
calcAccelsWithBoxPA epsilon llx lly rux ruy mpts
= let mpts' = [: MP x y m | (x, y, m) <- fromPArrayP mpts :]
accs' = calcAccelsWithBox epsilon llx lly rux ruy mpts'
in toPArrayP accs'
{-# NOINLINE calcAccelsWithBoxPA #-}
-- | Given the extend of a bounding box containing all the points,
-- calculate the accelerations on all of them.
calcAccelsWithBox
:: Double
-> Double -> Double -> Double -> Double
-> [: MassPoint :]
-> [: Accel :]
calcAccelsWithBox epsilon llx lly rux ruy mspts
= accs
where accs = [: calcAccel epsilon m tree | m <- mspts :]
tree = buildTree (Box llx lly rux ruy) mspts
-- | Build the Barnes-Hut quadtree tree.
buildTree :: BoundingBox -> [: MassPoint :] -> BHTree
buildTree bb particles
| lengthP particles I.<= 1 = BHT s x y m emptyP
| otherwise = BHT s x y m subTrees
where (MP x y m) = calcCentroid particles
(boxes, splitPnts) = splitPoints bb particles
subTrees = [:buildTree bb' ps | (bb', ps) <- zipP boxes splitPnts:]
(Box llx lly rux ruy) = bb
sx = rux D.- llx
sy = ruy D.- lly
s = if sx D.< sy then sx else sy
-- | Split massPoints according to their locations in the quadrants.
splitPoints
:: BoundingBox
-> [: MassPoint :]
-> ([:BoundingBox:], [:[: MassPoint :]:])
splitPoints b@(Box llx lly rux ruy) particles
| noOfPoints I.<= 1 = (singletonP b, singletonP particles)
| otherwise
= unzipP [: (b,p) | (b,p) <- zipP boxes splitPars, lengthP p I.> 0:]
where noOfPoints = lengthP particles
lls = [: p | p <- particles, inBox b1 p :]
lus = [: p | p <- particles, inBox b2 p :]
rus = [: p | p <- particles, inBox b3 p :]
rls = [: p | p <- particles, inBox b4 p :]
b1 = Box llx lly midx midy
b2 = Box llx midy midx ruy
b3 = Box midx midy rux ruy
b4 = Box midx lly rux midy
boxes = singletonP b1 +:+ singletonP b2 +:+ singletonP b3 +:+ singletonP b4
splitPars = singletonP lls +:+ singletonP lus +:+ singletonP rus +:+ singletonP rls
(midx, midy) = ((llx D.+ rux) D./ 2.0 , (lly D.+ ruy) D./ 2.0)
-- | Checks if particle is in box (excluding left and lower border)
inBox :: BoundingBox -> MassPoint -> Bool
inBox (Box llx lly rux ruy) (MP px py _)
= (px D.> llx) && (px D.<= rux) && (py D.> lly) && (py D.<= ruy)
-- | Calculate the centroid of some points.
calcCentroid:: [:MassPoint:] -> MassPoint
calcCentroid mpts
= MP (sumP xs / mass) (sumP ys / mass) mass
where mass = sumP [: m | MP _ _ m <- mpts :]
(xs, ys) = unzipP [: (m D.* x, m D.* y) | MP x y m <- mpts :]
-- | Calculate the accelleration of a point due to the points in the given tree.
calcAccel :: Double -> MassPoint -> BHTree -> (Double, Double)
calcAccel epsilon mpt (BHT s x y m subtrees)
| lengthP subtrees I.== 0
= accel epsilon mpt (MP x y m)
| isFar mpt s x y
= accel epsilon mpt (MP x y m)
| otherwise
= let (xs, ys) = unzipP [: calcAccel epsilon mpt st | st <- subtrees :]
in (sumP xs, sumP ys)
-- | Calculate the acceleration on a point due to some other point.
accel :: Double -- ^ If the distance between the points is smaller than this
-- then ignore the forces between them.
-> MassPoint -- ^ The point being acclerated.
-> MassPoint -- ^ Neibouring point.
-> Accel
accel epsilon (MP x1 y1 _) (MP x2 y2 m)
= (aabs D.* dx D./ r , aabs D.* dy D./ r)
where rsqr = (dx D.* dx) D.+ (dy D.* dy) D.+ epsilon D.* epsilon
r = sqrt rsqr
dx = x1 D.- x2
dy = y1 D.- y2
aabs = m D./ rsqr
-- | If the point is far from a cell in the tree then we can use
-- it's centroid as an approximation of all the points in the region.
isFar :: MassPoint -- point being accelerated
-> Double -- size of region
-> Double -- position of center of mass of cell
-> Double -- position of center of mass of cell
-> Bool
isFar (MP x1 y1 m) s x2 y2
= let dx = x2 D.- x1
dy = y2 D.- y1
dist = sqrt (dx D.* dx D.+ dy D.* dy)
in (s D./ dist) D.< 1
|
