<commit_before>{-# LANGUAGE NoImplicitPrelude, TypeFamilies #-}
{-# OPTIONS -Wall #-}

module Main(main) where

import           Data.Typeable
import           Language.Paraiso.Generator.Cpp
import           Language.Paraiso.OM.Builder
import           Language.Paraiso.OM.Builder.Boolean
import           Language.Paraiso.OM.DynValue as DVal
import           Language.Paraiso.OM.Graph
import           Language.Paraiso.OM.Realm 
import qualified Language.Paraiso.OM.Reduce as Reduce
import           Language.Paraiso.POM
import           Language.Paraiso.Prelude 
import           Language.Paraiso.Tensor
import           Hydro
import           System.Directory (createDirectoryIfMissing)

realDV :: DynValue
realDV = DynValue{DVal.realm = Local, DVal.typeRep = typeOf (0::Real)}

intGDV :: DynValue
intGDV = DynValue{DVal.realm = Global, DVal.typeRep = typeOf (0::Int)}

realGDV :: DynValue
realGDV = DynValue{DVal.realm = Global, DVal.typeRep = typeOf (0::Real)}


-- the list of static variables for this machine
pomSetup :: Setup Dim Int 
pomSetup = Setup $ 
            [Named (Name "generation") intGDV] ++
            [Named (Name "time") realGDV] ++
            [Named (Name "cfl") realGDV] ++
            foldMap (\name0 -> [Named name0 realGDV]) dRNames ++ 
            foldMap (\name0 -> [Named name0 realGDV]) extentNames ++ 
            [Named (Name "density") realDV]  ++
            foldMap (\name0 -> [Named name0 realDV]) velocityNames ++ 
            [Named (Name "pressure") realDV]  

velocityNames :: Dim (Name)
velocityNames = compose (\axis -> Name $ "velocity" ++ show (axisIndex axis))

dRNames :: Dim (Name)
dRNames = compose (\axis -> Name $ "dR" ++ show (axisIndex axis))

extentNames :: Dim (Name)
extentNames = compose (\axis -> Name $ "extent" ++ show (axisIndex axis))

loadReal :: Name -> BR
loadReal = load TLocal (undefined::Real) 
loadGReal :: Name -> BGR
loadGReal = load TGlobal (undefined::Real) 

----------------------------------------------------------------
-- Hydro utility functions.
----------------------------------------------------------------

hllc :: Axis Dim -> Hydro BR -> Hydro BR -> B (Hydro BR)
hllc i left right = do
  densMid  <- bind $ (density left    + density right   ) / 2
  soundMid <- bind $ (soundSpeed left + soundSpeed right) / 2
  let 
      speedLeft  = velocity left  !i
      speedRight = velocity right !i
  presStar <- bind $ max 0 $ (pressure left   + pressure right  ) / 2 -
              densMid * soundMid * (speedRight - speedLeft) 
  shockLeft <- bind $ velocity left !i - 
               soundSpeed left * hllcQ presStar (pressure left)
  shockRight <- bind $ velocity right !i +
               soundSpeed right * hllcQ presStar (pressure right)
  shockStar <- bind $ (pressure right - pressure left
		       + density left  * speedLeft  * (shockLeft  - speedLeft)
		       - density right * speedRight * (shockRight - speedRight) )
               / (density left  * (shockLeft  - speedLeft ) - 
                  density right * (shockRight - speedRight) )
  lesta <- starState shockStar shockLeft  left
  rista <- starState shockStar shockRight right
  let selector a b c d =
          select (0 `lt` shockLeft) a $ 
          select (0 `lt` shockStar) b $
          select (0 `lt` shockRight) c d
  mapM bind $ selector <$> left <*> lesta <*> rista <*> right
    where
      hllcQ sp p = select (p `le` sp) 1 $
                   sqrt(1 + (kGamma + 1)/(2*kGamma) * (sp / p - 1))
      starState starShock shock x = do
        let speed = velocity x !i
        dens <- bind $ density x * (shock - speed) / (shock - starShock)
        mome <- sequence $ compose 
                (\j ->  bind $ dens * if i==j then starShock
                                      else velocity x !j)
        enrg <- bind $ dens * 
                (energy x / density x + 
                 (starShock - speed) * 
                 (starShock + pressure x/density x/(shock - speed)))
        bindConserved dens mome enrg

buildProceed :: Builder Dim Int ()
buildProceed = do
  dens    <- bind $ loadReal $ Name "density"
  velo    <- mapM (bind . loadReal) velocityNames
  pres    <- bind $ loadReal $ Name "pressure"
  
  timeG   <- bind $ loadGReal $ Name "time"
  cflG    <- bind $ loadGReal $ Name "cfl"
  
  dRG     <- mapM (bind . loadGReal) dRNames  
  dR      <- mapM (bind . broadcast) dRG 
  cell <-  bindPrimitive dens velo pres
  
  let timescale i = dR!i / (soundSpeed cell + abs (velocity cell !i))
  dts <- bind $ foldl1 min $ compose timescale
  
  dtG <- bind $ cflG * reduce Reduce.Min dts
  dt  <- bind $ broadcast dtG
  
  cell2 <- proceedSingle 1 (dt/2) dR cell  cell
  cell3 <- proceedSingle 2  dt    dR cell2 cell
  
  store (Name "time") $ timeG + dtG
  store (Name "density") $ density cell3
  _ <- sequence $ compose(\i ->  store (velocityNames!i) $ velocity cell3 !i)
  store (Name "pressure") $ pressure cell3


proceedSingle :: Int -> BR -> Dim BR -> Hydro BR -> Hydro BR -> B (Hydro BR)
proceedSingle order dt dR cellF cellS = do
  let calcWall i = do
        (lp,rp) <- interpolate order i cellF
        hllc i lp rp
  wall <- sequence $ compose calcWall
  cx <- addFlux dt dR wall (Axis 0) cellS
  addFlux dt dR wall (Axis 1) cx
  
addFlux :: BR -> Dim BR -> Dim (Hydro BR) -> Axis Dim -> Hydro BR -> B (Hydro BR)
addFlux dt dR wall ex cell = do
  dtdx <- bind $ dt / dR!ex
  leftWall  <- mapM bind $ wall ! ex
  rightWall <- mapM (bind . shift (negate $ unitVector ex)) $ wall!ex
  dens1 <- bind $ density cell + dtdx * (densityFlux leftWall - densityFlux rightWall)!ex
  mome1 <- sequence $ compose 
           (\j -> bind $ (momentum cell !j + dtdx * 
                          (momentumFlux leftWall - momentumFlux rightWall) !j!ex))
  enrg1 <- bind $ energy  cell + dtdx * (energyFlux leftWall - energyFlux rightWall)  !ex
  
  bindConserved dens1 mome1 enrg1

interpolateSingle :: Int -> BR -> BR -> BR -> BR -> B (BR,BR)
interpolateSingle order x0 x1 x2 x3 = 
  if order == 1 
  then do
    return (x1, x2)<commit_msg>dimension dependent code eliminated.<commit_after>{-# LANGUAGE NoImplicitPrelude, TypeFamilies #-}
{-# OPTIONS -Wall #-}

module Main(main) where

import           Data.Typeable
import           Language.Paraiso.Generator.Cpp
import           Language.Paraiso.OM.Builder
import           Language.Paraiso.OM.Builder.Boolean
import           Language.Paraiso.OM.DynValue as DVal
import           Language.Paraiso.OM.Graph
import           Language.Paraiso.OM.Realm 
import qualified Language.Paraiso.OM.Reduce as Reduce
import           Language.Paraiso.POM
import           Language.Paraiso.Prelude 
import           Language.Paraiso.Tensor
import           Hydro
import           System.Directory (createDirectoryIfMissing)

realDV :: DynValue
realDV = DynValue{DVal.realm = Local, DVal.typeRep = typeOf (0::Real)}

intGDV :: DynValue
intGDV = DynValue{DVal.realm = Global, DVal.typeRep = typeOf (0::Int)}

realGDV :: DynValue
realGDV = DynValue{DVal.realm = Global, DVal.typeRep = typeOf (0::Real)}


-- the list of static variables for this machine
pomSetup :: Setup Dim Int 
pomSetup = Setup $ 
            [Named (Name "generation") intGDV] ++
            [Named (Name "time") realGDV] ++
            [Named (Name "cfl") realGDV] ++
            foldMap (\name0 -> [Named name0 realGDV]) dRNames ++ 
            foldMap (\name0 -> [Named name0 realGDV]) extentNames ++ 
            [Named (Name "density") realDV]  ++
            foldMap (\name0 -> [Named name0 realDV]) velocityNames ++ 
            [Named (Name "pressure") realDV]  

velocityNames :: Dim (Name)
velocityNames = compose (\axis -> Name $ "velocity" ++ show (axisIndex axis))

dRNames :: Dim (Name)
dRNames = compose (\axis -> Name $ "dR" ++ show (axisIndex axis))

extentNames :: Dim (Name)
extentNames = compose (\axis -> Name $ "extent" ++ show (axisIndex axis))

loadReal :: Name -> BR
loadReal = load TLocal (undefined::Real) 
loadGReal :: Name -> BGR
loadGReal = load TGlobal (undefined::Real) 

----------------------------------------------------------------
-- Hydro utility functions.
----------------------------------------------------------------

hllc :: Axis Dim -> Hydro BR -> Hydro BR -> B (Hydro BR)
hllc i left right = do
  densMid  <- bind $ (density left    + density right   ) / 2
  soundMid <- bind $ (soundSpeed left + soundSpeed right) / 2
  let 
      speedLeft  = velocity left  !i
      speedRight = velocity right !i
  presStar <- bind $ max 0 $ (pressure left   + pressure right  ) / 2 -
              densMid * soundMid * (speedRight - speedLeft) 
  shockLeft <- bind $ velocity left !i - 
               soundSpeed left * hllcQ presStar (pressure left)
  shockRight <- bind $ velocity right !i +
               soundSpeed right * hllcQ presStar (pressure right)
  shockStar <- bind $ (pressure right - pressure left
		       + density left  * speedLeft  * (shockLeft  - speedLeft)
		       - density right * speedRight * (shockRight - speedRight) )
               / (density left  * (shockLeft  - speedLeft ) - 
                  density right * (shockRight - speedRight) )
  lesta <- starState shockStar shockLeft  left
  rista <- starState shockStar shockRight right
  let selector a b c d =
          select (0 `lt` shockLeft) a $ 
          select (0 `lt` shockStar) b $
          select (0 `lt` shockRight) c d
  mapM bind $ selector <$> left <*> lesta <*> rista <*> right
    where
      hllcQ sp p = select (p `le` sp) 1 $
                   sqrt(1 + (kGamma + 1)/(2*kGamma) * (sp / p - 1))
      starState starShock shock x = do
        let speed = velocity x !i
        dens <- bind $ density x * (shock - speed) / (shock - starShock)
        mome <- sequence $ compose 
                (\j ->  bind $ dens * if i==j then starShock
                                      else velocity x !j)
        enrg <- bind $ dens * 
                (energy x / density x + 
                 (starShock - speed) * 
                 (starShock + pressure x/density x/(shock - speed)))
        bindConserved dens mome enrg

buildProceed :: Builder Dim Int ()
buildProceed = do
  dens    <- bind $ loadReal $ Name "density"
  velo    <- mapM (bind . loadReal) velocityNames
  pres    <- bind $ loadReal $ Name "pressure"
  
  timeG   <- bind $ loadGReal $ Name "time"
  cflG    <- bind $ loadGReal $ Name "cfl"
  
  dRG     <- mapM (bind . loadGReal) dRNames  
  dR      <- mapM (bind . broadcast) dRG 
  cell <-  bindPrimitive dens velo pres
  
  let timescale i = dR!i / (soundSpeed cell + abs (velocity cell !i))
  dts <- bind $ foldl1 min $ compose timescale
  
  dtG <- bind $ cflG * reduce Reduce.Min dts
  dt  <- bind $ broadcast dtG
  
  cell2 <- proceedSingle 1 (dt/2) dR cell  cell
  cell3 <- proceedSingle 2  dt    dR cell2 cell
  
  store (Name "time") $ timeG + dtG
  store (Name "density") $ density cell3
  _ <- sequence $ compose(\i ->  store (velocityNames!i) $ velocity cell3 !i)
  store (Name "pressure") $ pressure cell3


proceedSingle :: Int -> BR -> Dim BR -> Hydro BR -> Hydro BR -> B (Hydro BR)
proceedSingle order dt dR cellF cellS = do
  let calcWall i = do
        (lp,rp) <- interpolate order i cellF
        hllc i lp rp
  wall <- sequence $ compose calcWall
  foldl1 (.) (compose (\i -> (>>= addFlux dt dR wall i))) $ return cellS
  
--  cx <- addFlux dt dR wall (Axis 0) cellS
--  addFlux dt dR wall (Axis 1) cx
  
addFlux :: BR -> Dim BR -> Dim (Hydro BR) -> Axis Dim -> Hydro BR -> B (Hydro BR)
addFlux dt dR wall ex cell = do
  dtdx <- bind $ dt / dR!ex
  leftWall  <- mapM bind $ wall ! ex
  rightWall <- mapM (bind . shift (negate $ unitVector ex)) $ wall!ex
  dens1 <- bind $ density cell + dtdx * (densityFlux leftWall - densityFlux rightWall)!ex
  mome1 <- sequence $ compose 
           (\j -> bind $ (momentum cell !j + dtdx * 
                          (momentumFlux leftWall - momentumFlux rightWall) !j!ex))
  enrg1 <- bind $ energy  cell + dtdx * (energyFlux leftWall - energyFlux rightWall)  !ex
  
  bindConserved dens1 mome1 enrg1

interpolateSingle :: Int -> BR -> BR -> BR -> BR -> B (BR,BR)
interpolateSingle order x0 x1 x2 x3 = 
  if order == 1 
  then do
    return (x1, x2)
{'fraction_non_alphanumeric': 0.13017702758336763, 'fraction_numerical': 0.007657472210786332, 'mean_word_length': 2.100842481490937, 'pattern_counts': {'":': 0, '<': 79, '<?xml version=': 0, '>': 72, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2459247', 'n_tokens_mistral': 4006, 'n_tokens_neox': 3717, 'n_words': 1260}