cl3-3.0.0.0: benchmarks/NbodyMassiv.hs

{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE InstanceSigs #-}



import Algebra.Geometric.Cl3
import Data.Massiv.Array as A
import Control.Monad (replicateM_,when)
import Data.Time.Clock.POSIX (getPOSIXTime)

import Foreign.Ptr (Ptr, castPtr, plusPtr)
import Foreign.Storable (Storable(..))

import Data.IORef (IORef(),newIORef,modifyIORef',readIORef)


main :: IO ()
main = do
  arrayProgrammingAlgorithm
  bitTwiddlingAlgorithm
  
arrayProgrammingAlgorithm :: IO ()
arrayProgrammingAlgorithm = 
  do tStart <- getPOSIXTime
     n <- return (50000000 :: Int)
     planetsPos :: Array S Ix1 Cl3_V3 <- fromListsM Seq planetspos  -- generate immutible Massiv array of the planets position
     planetsVel :: Array S Ix1 Cl3_V3 <- fromListsM Seq planetsvel  -- generate immutible Massiv array of the planets velocity
     planetsMass :: Array S Ix1 Cl3_R <- fromListsM Seq planetsmass  -- generate immutible Massiv array of the planets mass
     mutablePos <- thaw planetsPos  -- copy to a mutible Massiv array of the planets position
     mutableVel <- thaw planetsVel  -- copy to a mutible Massiv array of the planets velocity
     let (Sz1 c) = size planetsVel  -- Size of the velocity array
         c' = (c^(2 :: Int) - c) `div` 2  -- size of the nChoose2With array
     mutableScratch <- newMArray (Sz1 c') (toCl3_V3 $ V3 0 0 0) -- nChoose2With sized scratch pad
     mutableScratch2D <- newMArray (Sz (c :. (c-1))) (toCl3_V3 $ V3 0 0 0) -- 2D scratch pad for holding acceleration components
     offsetMomentum mutableVel planetsMass -- modify the Sun's velocity such that the total momentum is zero
     printEnergy mutablePos mutableVel planetsMass
     replicateM_ n (advance mutablePos mutableVel planetsMass mutableScratch mutableScratch2D)
     printEnergy mutablePos mutableVel planetsMass
     tEnd <- getPOSIXTime
     print $ show (tEnd - tStart) ++ " (sec) for " ++ show n ++ " iterations." -- Relative time measurement

-- | 'bitTwiddlingAlgorithm' with Ϟ symbolizing the bit twiddling algorithm versions
bitTwiddlingAlgorithm :: IO ()
bitTwiddlingAlgorithm = 
  do tStart <- getPOSIXTime
     n <- return (50000000 :: Int)
     (solarSystem,sz) <- offsetMomentumϞ [sun,jupiter,saturn,uranus,neptune] -- returns a zero momentum solar system
     printEnergyϞ solarSystem sz
     replicateM_ n (advanceϞ solarSystem sz)
     printEnergyϞ solarSystem sz
     tEnd <- getPOSIXTime
     print $ show (tEnd - tStart) ++ " (sec) for " ++ show n ++ " iterations." -- Relative time measurement

offsetMomentumϞ :: [Body] -> IO (SolarSystem,Int) -- A SolarSystem and the Size : size: sz :: Int
offsetMomentumϞ bodies = do
  initBodies :: Array S Ix1 Body <- A.fromListsM Seq bodies
  -- mutate the solar system so that the sun's velocity gives zero total momentum
  let velocity :: Array D Ix1 Cl3 = A.map (fromCl3_V3.vel) initBodies
      mass :: Array D Ix1 Cl3 = A.map (fromCl3_R.mas) initBodies
  ss :: SolarSystem <- thaw initBodies
  let sz = totalElem.sizeOfMArray $ ss
  bs <- readM ss 0
  let sunV = fromCl3_V3.vel $ bs
      sunM = fromCl3_R.mas $ bs
  writeM ss 0 (bs{vel = toCl3_V3 (sunV - ((velocity `dot_V3` mass) / sunM))})
  return (ss,sz)


printEnergyϞ :: SolarSystem -> Int -> IO ()
printEnergyϞ ss sz = do
  acc :: IORef Cl3 <- newIORef (R 0)
  ssFroz <- freeze Seq ss
  iforIO_ ssFroz $ \idx bi ->
    do let bi_pos :: Cl3 = fromCl3_V3.pos $ bi
           bi_vel :: Cl3 = fromCl3_V3.vel $ bi
           bi_mas :: Cl3 = fromCl3_R.mas $ bi
           ke :: Cl3 = 0.5 * bi_mas * (toR $ bi_vel^(2::Int))
       modifyIORef' acc (+ ke)  -- accumulate the kenetic energy for each body
       -- now calcualate the potential energy
       let szMinus1 = sz-1
       when (idx < szMinus1) $ do
         brest <- extractM (idx+1) (Sz1 (szMinus1-idx)) ssFroz
         forIO_ brest $ \bj ->
           do let bj_pos :: Cl3 = fromCl3_V3.pos $ bj
                  bj_mas :: Cl3 = fromCl3_R.mas $ bj
                  distance :: Cl3 = sqrt.toR $ (bi_pos - bj_pos)^(2::Int)
                  pe_ij :: Cl3 = (bi_mas * bj_mas) / distance
              modifyIORef' acc (\ac -> ac - pe_ij)
  energy <- readIORef acc
  print energy

advanceϞ :: SolarSystem -> Int -> IO ()
advanceϞ ss sz = do
  updateVelocityϞ ss sz
  updatePositionϞ ss

updateVelocityϞ :: SolarSystem -> Int -> IO ()
updateVelocityϞ ss sz = do
  ssFroz <- freeze Seq ss
  iforIO_ ssFroz $ \idx bi ->
    when (idx < sz-1) $ do
      let bi_pos :: Cl3 = fromCl3_V3.pos $ bi
          bi_mas :: Cl3 = fromCl3_R.mas $ bi
      brest <- extractM (idx+1) (Sz1 (sz-idx-1)) ssFroz
      iforIO_ brest $ \jdx bj ->
        do let bj_pos :: Cl3 = fromCl3_V3.pos $ bj
               bj_mas :: Cl3 = fromCl3_R.mas $ bj
               (absMag,dposUnit) = abssignum $ bi_pos - bj_pos
               mag = dt / ((absMag)^(2::Int))
               aij = negate $ dposUnit * bj_mas * mag
               aji = dposUnit * bi_mas * mag
           modify_ ss (modVel aij) idx
           modify_ ss (modVel aji) (idx+jdx+1)

modVel :: PrimMonad m => Cl3 -> Body -> m Body
modVel a bod = do
  let vbod = fromCl3_V3 $ vel bod
  return bod{vel = toCl3_V3 $ vbod + a}

updatePositionϞ :: SolarSystem -> IO ()
updatePositionϞ ss = do
  iforPrimM_ ss $ \idx b ->
    do let p :: Cl3 = fromCl3_V3.pos $ b
           v :: Cl3 = fromCl3_V3.vel $ b
       writeM ss idx b{pos = toCl3_V3 (p + dt * v)}
        

advance :: MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix1 Cl3_V3 -> Array S Ix1 Cl3_R -> MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix2 Cl3_V3 -> IO ()
advance sysPos sysVel sysMass scratchPad scratchPad2D =
  do updateVelocity sysPos sysVel sysMass scratchPad scratchPad2D
     updatePosition sysPos sysVel

updateVelocity :: MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix1 Cl3_V3 -> Array S Ix1 Cl3_R -> MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix2 Cl3_V3 -> IO ()
updateVelocity sysPos sysVel sysMass scratchPad scratchPad2D =
  do -- Calculate the inverse square distance and unit vector between the bodies
     frozPos <- freeze Seq sysPos
     let ps :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozPos
         scaledPij = nChoose2With (\ i j -> let (a,s) = abssignum $ i - j in (recip (a^(2 :: Int))) * s) ps -- probably should use a computeSource here, because scaledPij is used twice, or verify streams are still working.
     computeInto scratchPad (A.map toCl3_V3 scaledPij)  -- Compute result to the scratch pad
     frozScratch <- freeze Seq scratchPad  -- Freeze the result for future use
     -- Scale the previous inverse square distance by the mass of the other body
     let scratch :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozScratch
         ms :: Array D Ix1 Cl3 = A.map fromCl3_R sysMass
         upper = nChoose2With (\ _ j -> negate j) ms !*! scratch
         lower = nChoose2With (\ i _ -> i) ms !*! scratch
     computeInto scratchPad (A.map toCl3_V3 upper)
     frozScratchUpper <- freeze Seq scratchPad
     computeInto scratchPad (A.map toCl3_V3 lower)
     frozScratchLower <- freeze Seq scratchPad
     -- Load each body acceleration contribution into a 2D array 
     frozVel <- freeze Seq sysVel
     let vs :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozVel
         rows = elemsCount vs
     iforPrimM_ scratchPad2D (\ix _ -> write_ scratchPad2D ix (genAcc rows frozScratchUpper frozScratchLower ix))
     frozScratch2D <- freeze Seq scratchPad2D
     --  Sum each row to calcuate the total acceleration and then update the velocity
     let accComponents :: Array D Ix2 Cl3 = A.map fromCl3_V3 frozScratch2D
         as = makeArray Seq (size vs) (\ ix -> foldlS (+) (V3 0 0 0) (accComponents !> ix))
         vs' = vs !+! (dt *. as)
     computeInto sysVel (A.map toCl3_V3 vs')

genAcc :: Int -> Array S Ix1 Cl3_V3 -> Array S Ix1 Cl3_V3 -> Ix2 -> Cl3_V3
genAcc rows upper lower (row :. col) | row <= col = upper ! unSplitIdx rows (row,col+1)
                                     | otherwise  = lower ! unSplitIdx rows (col,row)

unSplitIdx :: Int -> (Int,Int) -> Int
unSplitIdx c (row,col) = go 0 (c-1) 0 + (col-1 - row)
  where
    go n c1 acc | n == row = acc
                | otherwise = go (n+1) (c1-1) (acc+c1)



updatePosition :: MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix1 Cl3_V3 -> IO ()
updatePosition sysPos sysVel =
  do frozPos <- freeze Seq sysPos
     frozVel <- freeze Seq sysVel
     let ps :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozPos
         vs :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozVel
         ps' = ps !+! (dt *. vs)
     computeInto sysPos (A.map toCl3_V3 ps')


offsetMomentum :: MArray RealWorld S Ix1 Cl3_V3 -> Array S Ix1 Cl3_R -> IO ()
offsetMomentum sysVel sysMass =
  do frozVel <- freeze Seq sysVel
     planetsvs :: Array D Ix1 Cl3 <- extractM 1 4 (A.map fromCl3_V3 frozVel)  -- select just the planets and not the sun
     planetms :: Array D Ix1 Cl3 <- extractM 1 4 (A.map fromCl3_R sysMass)  -- select just the planets and not the sun
     (fromCl3_R -> sunMassCl3) <- sysMass !? 0
     writeM sysVel 0 (toCl3_V3 $ negate (planetsvs `dot_V3` planetms) / sunMassCl3)
     

-- dot_V3 is to work around the typical sum in the dot product that has an aculumulator of a Real 0 value
dot_V3 :: Array D Ix1 Cl3 -> Array D Ix1 Cl3 -> Cl3
dot_V3 v1 v2 = foldlS (+) (V3 0 0 0) (A.zipWith (*) v1 v2)


printEnergy :: MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix1 Cl3_V3 -> Array S Ix1 Cl3_R -> IO ()
printEnergy sysPos sysVel sysMass =
  do frozPos <- freeze Seq sysPos
     frozVel <- freeze Seq sysVel
     let ps :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozPos
         vs :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozVel
         ms :: Array D Ix1 Cl3 = A.map fromCl3_R sysMass
         ke = 0.5 * (ms !.! A.map toR (vs.^2)) -- Kenetic Energy O(n)
         pe = negate.A.sum $ nChoose2With (*) ms !/! nChoose2With (\ x y -> abs $ x - y) ps -- Gravitational Potential Energy O((n^2 - n)/2)
     print $ ke + pe -- print the total energy

-- This function is a bit of a combinatorial explosion so it is marked to be evaluated in Seqallel, probably should convert to a sgenerate :: Sz1 -> (Ix1 -> e) -> Vector DS e 
nChoose2With :: (Cl3 -> Cl3 -> Cl3) -> Array D Ix1 Cl3 -> Array D Ix1 Cl3
nChoose2With f bs =
  let (Sz1 c) = size bs
      c' = (c^(2 :: Int) - c) `div` 2
      bs' :: Array BN Ix1 Cl3 = computeSource bs
  in makeArray Seq (Sz1 c') (genElements f bs' c)

genElements :: (Cl3 -> Cl3 -> Cl3) -> Array BN Ix1 Cl3 -> Int -> Ix1 -> Cl3
genElements f bs c idx =
  let (ix1,ix2) = splitIdx c idx
  in f (bs ! ix1) (bs ! ix2)

splitIdx :: Int -> Int -> (Int,Int)
splitIdx c = go (0,1) (c-2)
  where
    go (row,col) x ix | ix - x <= 0 = (row, col + ix)
                      | otherwise = go (row+1,col+1) (x-1) (ix-x-1)

planetspos :: [Cl3_V3]
planetspos = pos <$> [sun,jupiter,saturn,uranus,neptune]

planetsvel :: [Cl3_V3]
planetsvel = vel <$> [sun,jupiter,saturn,uranus,neptune]

planetsmass :: [Cl3_R]
planetsmass = mas <$> [sun,jupiter,saturn,uranus,neptune]

type SolarSystem = MArray RealWorld S Ix1 Body

data Body = Body
  {pos :: Cl3_V3
  ,vel :: Cl3_V3
  ,mas :: Cl3_R}

instance Storable Body where
  sizeOf _ = sizeOf (undefined :: Cl3_V3) + sizeOf (undefined :: Cl3_V3) + sizeOf (undefined :: Cl3_R)
  alignment _ = sizeOf (undefined :: Double) -- alignment must be a 2^x
  peekElemOff ptr idx = peek (castPtr ptr `plusPtr` (idx * sizeOf (undefined :: Body)))
  pokeElemOff ptr idx e = poke (castPtr ptr `plusPtr` (idx * sizeOf (undefined :: Body))) e
  peek :: Ptr Body -> IO Body
  peek ptr = do
    p :: Cl3_V3 <- peek (castPtr ptr `plusPtr` 0)
    v :: Cl3_V3 <- peek (castPtr ptr `plusPtr` (sizeOf p))
    m :: Cl3_R <- peek (castPtr ptr `plusPtr` (sizeOf p) `plusPtr` (sizeOf v))
    return (Body p v m)
  poke :: Ptr Body -> Body -> IO ()
  poke ptr (Body p v m) = do
    poke (castPtr ptr `plusPtr` 0) p
    poke (castPtr ptr `plusPtr` (sizeOf p)) v
    poke (castPtr ptr `plusPtr` (sizeOf p) `plusPtr` (sizeOf v)) m

sun :: Body
sun = Body (toCl3_V3 $ V3 0 0 0) -- Sun Position
           (toCl3_V3 $ V3 0 0 0) -- Sun Velocity 
           (toCl3_R $ R solarMass) -- Sun Mass
jupiter :: Body
jupiter = Body (toCl3_V3 $ V3 (4.84143144246472090e+00) (-1.16032004402742839e+00) (-1.03622044471123109e-01))
               (toCl3_V3 $ V3 (1.66007664274403694e-03*dp) (7.69901118419740425e-03*dp) (-6.90460016972063023e-05*dp))
               (toCl3_R $ R (9.54791938424326609e-04 * solarMass))
saturn :: Body
saturn = Body (toCl3_V3 $ V3 (8.34336671824457987e+00) (4.12479856412430479e+00) (-4.03523417114321381e-01))
              (toCl3_V3 $ V3 (-2.76742510726862411e-03*dp) (4.99852801234917238e-03*dp) (2.30417297573763929e-05*dp))
              (toCl3_R $ R (2.85885980666130812e-04 * solarMass))
uranus :: Body
uranus = Body (toCl3_V3 $ V3 (1.28943695621391310e+01) (-1.51111514016986312e+01) (-2.23307578892655734e-01))
              (toCl3_V3 $ V3 (2.96460137564761618e-03*dp) (2.37847173959480950e-03*dp) (-2.96589568540237556e-05*dp))
              (toCl3_R $ R (4.36624404335156298e-05 * solarMass))
neptune :: Body
neptune = Body (toCl3_V3 $ V3 (1.53796971148509165e+01) (-2.59193146099879641e+01) (1.79258772950371181e-01))
               (toCl3_V3 $ V3 (2.68067772490389322e-03*dp) (1.62824170038242295e-03*dp) (-9.51592254519715870e-05*dp))
               (toCl3_R $ R (5.15138902046611451e-05 * solarMass))


daysPerYear :: Double
daysPerYear = 365.24

solarMass :: Double
solarMass = 4 * pi^(2 :: Int)
dp :: Double
dp = daysPerYear
dt :: Cl3
dt = R 0.01