{-# LANGUAGE BangPatterns, PatternGuards #-}
-- | The list version of the solver also builds the bounding box at every
-- node of the tree, which is good for visualisation.
module Solver.ListBH.Solver
( MassPoint (..)
, BoundingBox (..)
, BHTree (..)
, calcAccels
, buildTree)
where
import Common.Body
eClose :: Double
eClose = square 500
square x = x * x
-- | A rectangular region in 2D space.
data BoundingBox
= Box
{ boxLowerLeftX :: {-# UNPACK #-} !Double
, boxLowerLeftY :: {-# UNPACK #-} !Double
, boxUpperRightX :: {-# UNPACK #-} !Double
, boxUpperRightY :: {-# UNPACK #-} !Double }
deriving Show
-- | The Barnes-Hut tree we use to organise the points.
data BHTree
= BHT
{ bhTreeBox :: {-# UNPACK #-} !BoundingBox
, bhTreeCenterX :: {-# UNPACK #-} !Double
, bhTreeCenterY :: {-# UNPACK #-} !Double
, bhTreeMass :: {-# UNPACK #-} !Double
, bhTreeBranch :: ![BHTree] }
deriving Show
-- | Compute the acclerations on all these points.
calcAccels :: Double -> [MassPoint] -> [Accel]
calcAccels epsilon mpts
= map (calcAccel epsilon (buildTree mpts)) mpts
-- | Build a Barnes-Hut tree from these points.
buildTree :: [MassPoint] -> BHTree
buildTree mpts
= let (llx, lly, rux, ruy) = findBounds mpts
box = Box llx lly rux ruy
in buildTreeWithBox box mpts
-- | Find the coordinates of the bounding box that contains these points.
findBounds :: [MassPoint] -> (Double, Double, Double, Double)
{-# INLINE findBounds #-}
findBounds ((x1, y1, _) : rest1)
= go x1 y1 x1 y1 rest1
where go !left !right !down !up pts
= case pts of
[] -> (left, down, right, up)
(x, y, _) : rest
-> let left' = min left x
right' = max right x
down' = min down y
up' = max up y
in go left' right' down' up' rest
-- | Given a bounding box that contains all the points,
-- build the Barnes-Hut tree for them.
buildTreeWithBox
:: BoundingBox -- ^ bounding box containing all the points.
-> [MassPoint] -- ^ points in the box.
-> BHTree
buildTreeWithBox bb particles
| length particles <= 1 = BHT bb x y m []
| otherwise = BHT bb x y m subTrees
where (x, y, m) = calcCentroid particles
(boxes, splitPnts) = splitPoints bb particles
subTrees = [buildTreeWithBox bb' ps | (bb', ps) <- zip boxes splitPnts]
-- | Split massPoints according to their locations in the quadrants.
splitPoints
:: BoundingBox -- ^ bounding box containing all the points.
-> [MassPoint] -- ^ points in the box.
-> ( [BoundingBox] --
, [[MassPoint]])
splitPoints b@(Box llx lly rux ruy) particles
| noOfPoints <= 1 = ([b], [particles])
| otherwise
= unzip [ (b,p) | (b,p) <- zip boxes splitPars, length p > 0]
where
noOfPoints = length particles
-- The midpoint of the parent bounding box.
(midx, midy) = ((llx + rux) / 2.0 , (lly + ruy) / 2.0)
-- Split the parent bounding box into four quadrants.
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 = [b1, b2, b3, b4]
-- Sort the particles into the smaller boxes.
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 ]
splitPars = [lls, lus, rus, rls]
-- | Check if a particle is in box (excluding left and lower border)
inBox:: BoundingBox -> MassPoint -> Bool
{-# INLINE inBox #-}
inBox (Box llx lly rux ruy) (px, py, _)
= (px > llx) && (px <= rux) && (py > lly) && (py <= ruy)
-- | Calculate the centroid of some points.
calcCentroid :: [MassPoint] -> MassPoint
{-# INLINE calcCentroid #-}
calcCentroid mpts = (sum xs / mass, sum ys / mass, mass)
where
mass = sum [ m | (_, _, m) <- mpts ]
(xs, ys) = unzip [ (m * x, m * y) | (x, y, m) <- mpts ]
-- | Calculate the accelleration of a point due to the points in the given tree.
-- If the distance between the points is less then some small number
-- we set the accel to zero to avoid the acceleration going to infinity
-- and the points escaping the simulation.
--
-- We also use this behavior as a hacky way to discard the acceleration
-- of a point due to interaction with itself.
--
calcAccel:: Double -> BHTree -> MassPoint -> (Double, Double)
calcAccel !epsilon (BHT _ x y m subtrees) mpt
| [] <- subtrees
= accel epsilon mpt (x, y, m)
| not $ isClose mpt x y
= accel epsilon mpt (x, y, m)
| otherwise
= let (xs, ys) = unzip [ calcAccel epsilon st mpt | st <- subtrees]
in (sum xs, sum ys)
-- | If the a point is "close" to a region in the Barnes-Hut tree then we compute
-- the "real" acceleration on it due to all the points in the region, otherwise
-- we just use the centroid as an approximation of all the points in the region.
--
isClose :: MassPoint -> Double -> Double -> Bool
{-# INLINE isClose #-}
isClose (x1, y1, m) x2 y2
= (x1-x2) * (x1-x2) + (y1-y2) * (y1-y2) < eClose