{-# LANGUAGE ParallelArrays #-} {-# OPTIONS -fvectorise #-} module Solver.NestedBH.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 D./ mass) (sumP ys D./ 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