Changeset 31420 for lang/haskell

Show
Ignore:
Timestamp:
03/21/09 00:03:55 (4 years ago)
Author:
mokehehe
Message:

add type declaration

Files:
1 modified

Legend:

Unmodified
Added
Removed
  • lang/haskell/aobench/ao.hs

    r30638 r31420  
    1  
    21module Main () where 
    32 
     
    109 
    1110data Vec = Vec { 
    12         x :: Double, 
    13         y :: Double, 
    14         z :: Double 
    15         } deriving (Eq) 
     11      x :: Double, 
     12      y :: Double, 
     13      z :: Double 
     14    } deriving (Eq) 
    1615 
    1716data Isect = Isect { 
    18         isect_t :: Double, 
    19         isect_p :: Vec, 
    20         isect_n :: Vec 
    21         } deriving (Eq) 
     17      isect_t :: Double, 
     18      isect_p :: Vec, 
     19      isect_n :: Vec 
     20    } deriving (Eq) 
    2221 
    2322instance Ord Isect where 
    24         i1 < i2         = isect_t i1 < isect_t i2 
    25         i1 > i2         = isect_t i1 > isect_t i2 
    26         i1 <= i2        = isect_t i1 <= isect_t i2 
    27         i1 >= i2        = isect_t i1 >= isect_t i2 
     23    i1 < i2   = isect_t i1 < isect_t i2 
     24    i1 > i2   = isect_t i1 > isect_t i2 
     25    i1 <= i2  = isect_t i1 <= isect_t i2 
     26    i1 >= i2  = isect_t i1 >= isect_t i2 
    2827 
    2928data Sphere = Sphere { 
    30         sphere_center :: Vec, 
    31         sphere_radius :: Double 
    32         } 
     29      sphere_center :: Vec, 
     30      sphere_radius :: Double 
     31    } 
    3332 
    3433data Plane = Plane { 
    35         plane_p :: Vec, 
    36         plane_n :: Vec 
    37         } 
     34      plane_p :: Vec, 
     35      plane_n :: Vec 
     36    } 
    3837 
    3938data Ray = Ray { 
    40         ray_org :: Vec, 
    41         ray_dir :: Vec 
    42         } 
    43  
     39      ray_org :: Vec, 
     40      ray_dir :: Vec 
     41    } 
     42 
     43type Point = (Int, Int) 
     44 
     45-- Floatint-point Color 
     46type FColor = Vec 
     47 
     48sq :: Num a => a -> a 
    4449sq x = x * x 
    4550 
     51clamp :: Ord a => a -> a -> a -> a 
    4652clamp lower upper = min upper . max lower 
    4753 
     54i2d :: Int -> Double 
     55i2d = toEnum 
     56 
     57fpcol2icol vec = (f2i (x vec), f2i (y vec), f2i (z vec)) 
     58f2i = clamp 0 255 . floor . (* 255.99) 
     59 
     60drand48 :: IO Double 
    4861drand48 = getStdGen >>= random 
    4962 
    50 justIf cond val 
    51         | cond          = Just val 
    52         | otherwise     = Nothing 
     63justIf :: a -> Bool -> Maybe a 
     64val `justIf` cond 
     65    | cond      = Just val 
     66    | otherwise = Nothing 
    5367 
    5468zerovec = Vec { x = 0, y = 0, z = 0 } 
     69white = Vec { x = 1, y = 1, z = 1 } 
    5570 
    5671v0 `vadd` v1 = Vec (x v0 + x v1) (y v0 + y v1) (z v0 + z v1) 
     
    6176 
    6277v0 `vcross` v1 = Vec cx cy cz 
    63         where 
    64                 cx = y v0 * z v1 - z v0 * y v1 
    65                 cy = z v0 * x v1 - x v0 * z v1 
    66                 cz = x v0 * y v1 - y v0 * x v1 
     78    where 
     79      cx = y v0 * z v1 - z v0 * y v1 
     80      cy = z v0 * x v1 - x v0 * z v1 
     81      cz = x v0 * y v1 - y v0 * x v1 
    6782 
    6883vlength c = sqrt $ c `vdot` c 
    6984 
    7085vnormalize c 
    71         | l > 1.0e-17   = c `vscale` (1.0 / l) 
    72         | otherwise             = c 
    73         where 
    74                 l = vlength c 
    75  
     86    | l > 1.0e-17       = c `vscale` (1.0 / l) 
     87    | otherwise         = c 
     88    where 
     89      l = vlength c 
     90 
     91vaverage :: [Vec] -> Vec 
     92vaverage vecs = vsum `vscale` (1.0 / i2d vecnum) 
     93    where vsum = foldl vadd zerovec vecs 
     94          vecnum = length vecs 
     95 
     96ray_sphere_intersect :: Ray -> Sphere -> Maybe Isect 
    7697ray_sphere_intersect ray sphere = 
    77         justIf (d > 0.0 && t > 0.0) isect 
    78         where 
    79                 rs = ray_org ray `vsub` sphere_center sphere 
    80                 b = rs `vdot` ray_dir ray 
    81                 c = rs `vdot` rs - sq (sphere_radius sphere) 
    82                 d = b * b - c 
    83                 t = -b - sqrt d 
    84                 hitpos = ray_org ray `vadd` (ray_dir ray `vscale` t) 
    85                 isect = Isect { 
    86                         isect_t = t, 
    87                         isect_p = hitpos, 
    88                         isect_n = vnormalize $ hitpos `vsub` sphere_center sphere 
    89                         } 
    90  
     98    isect `justIf` (d > 0.0 && t > 0.0) 
     99    where 
     100      rs = ray_org ray `vsub` sphere_center sphere 
     101      b = rs `vdot` ray_dir ray 
     102      c = rs `vdot` rs - sq (sphere_radius sphere) 
     103      d = b * b - c 
     104      t = -b - sqrt d 
     105      hitpos = ray_org ray `vadd` (ray_dir ray `vscale` t) 
     106      isect = Isect { 
     107                isect_t = t, 
     108                isect_p = hitpos, 
     109                isect_n = vnormalize $ hitpos `vsub` sphere_center sphere 
     110              } 
     111 
     112ray_plane_intersect :: Ray -> Plane -> Maybe Isect 
    91113ray_plane_intersect ray plane = 
    92         justIf (abs v >= 1.0e-17 && t > 0) isect 
    93         where 
    94                 d = -(plane_p plane `vdot` plane_n plane) 
    95                 v = ray_dir ray `vdot` plane_n plane 
    96                 t = -(ray_org ray `vdot` plane_n plane + d) / v 
    97                 isect = Isect { 
    98                         isect_t = t, 
    99                         isect_p = ray_org ray `vadd` (ray_dir ray `vscale` t), 
    100                         isect_n = plane_n plane 
    101                         } 
    102  
     114    isect `justIf` (abs v >= 1.0e-17 && t > 0) 
     115    where 
     116      d = -(plane_p plane `vdot` plane_n plane) 
     117      v = ray_dir ray `vdot` plane_n plane 
     118      t = -(ray_org ray `vdot` plane_n plane + d) / v 
     119      isect = Isect { 
     120                isect_t = t, 
     121                isect_p = ray_org ray `vadd` (ray_dir ray `vscale` t), 
     122                isect_n = plane_n plane 
     123              } 
     124 
     125orthoBasis :: Vec -> (Vec, Vec, Vec) 
    103126orthoBasis n = (basis0, basis1, basis2) 
    104         where 
    105                 basis2 = n 
    106                 basis1' 
    107                         | x n < 0.6 && x n > -0.6       = Vec { x = 1, y = 0, z = 0 } 
    108                         | y n < 0.6 && y n > -0.6       = Vec { x = 0, y = 1, z = 0 } 
    109                         | z n < 0.6 && z n > -0.6       = Vec { x = 0, y = 0, z = 1 } 
    110                         | otherwise                                     = Vec { x = 1, y = 0, z = 0 } 
    111                 basis0 = vnormalize $ basis1' `vcross` basis2 
    112                 basis1 = vnormalize $ basis2 `vcross` basis0 
    113  
     127    where 
     128      basis2 = n 
     129      basis1' 
     130          | x n < 0.6 && x n > -0.6  = Vec { x = 1, y = 0, z = 0 } 
     131          | y n < 0.6 && y n > -0.6  = Vec { x = 0, y = 1, z = 0 } 
     132          | z n < 0.6 && z n > -0.6  = Vec { x = 0, y = 0, z = 1 } 
     133          | otherwise                = Vec { x = 1, y = 0, z = 0 } 
     134      basis0 = vnormalize $ basis1' `vcross` basis2 
     135      basis1 = vnormalize $ basis2 `vcross` basis0 
     136 
     137hemisphereDirectionRay :: Vec -> (Vec, Vec, Vec) -> Double -> Double -> Ray 
    114138hemisphereDirectionRay p (basis0, basis1, basis2) r1 r2 = 
    115         Ray { ray_org = p, ray_dir = r } 
    116         where 
    117                 theta = sqrt r1 
    118                 phi = 2.0 * pi * r2 
    119                 v = Vec { 
    120                         x = cos phi * theta, 
    121                         y = sin phi * theta, 
    122                         z = sqrt $ 1.0 - theta * theta 
    123                         } 
    124                 -- local -> global 
    125                 r = Vec { 
    126                         x = x v * x basis0 + y v * x basis1 + z v * x basis2, 
    127                         y = x v * y basis0 + y v * y basis1 + z v * y basis2, 
    128                         z = x v * z basis0 + y v * z basis1 + z v * z basis2 
    129                         } 
     139    Ray { ray_org = p, ray_dir = r } 
     140    where 
     141      theta = sqrt r1 
     142      phi = 2.0 * pi * r2 
     143      v = Vec { 
     144            x = cos phi * theta, 
     145            y = sin phi * theta, 
     146            z = sqrt $ 1.0 - sq theta 
     147          } 
     148      -- local -> global 
     149      r = Vec { 
     150            x = x v * x basis0 + y v * x basis1 + z v * x basis2, 
     151            y = x v * y basis0 + y v * y basis1 + z v * y basis2, 
     152            z = x v * z basis0 + y v * z basis1 + z v * z basis2 
     153          } 
     154 
     155randomHemisphereRay :: Vec -> Vec -> IO Ray 
     156randomHemisphereRay p n = do 
     157  r1 <- drand48 
     158  r2 <- drand48 
     159  return $ hemisphereDirectionRay p (orthoBasis n) r1 r2 
    130160 
    131161data Object = ObjSphere Sphere | ObjPlane Plane 
    132162 
     163findRayObjIntersect :: Ray -> Object -> Maybe Isect 
    133164findRayObjIntersect ray (ObjSphere sphere) = ray_sphere_intersect ray sphere 
    134165findRayObjIntersect ray (ObjPlane  plane)  = ray_plane_intersect  ray plane 
    135166 
     167isHitOccluder :: [Object] -> Ray -> Bool 
     168isHitOccluder objs ray = 
     169    not $ null $ catMaybes $ map (findRayObjIntersect ray) objs 
     170 
     171ambient_occlusion :: [Object] -> Int -> Isect -> IO FColor 
    136172ambient_occlusion objs nao_samples isect = do 
    137         randomRays <- mapM chooseRandomRay samples 
    138         let occludedNum = length $ filter id $ map isHitOccluder randomRays 
    139         return $ vscale white $ calcOccludedRatio occludedNum 
    140         where 
    141                 ntheta = nao_samples 
    142                 nphi = nao_samples 
    143                 eps = 0.0001 
    144                 p = isect_p isect `vadd` (isect_n isect `vscale` eps) 
    145                 basis = orthoBasis $ isect_n isect 
    146                 samples = [(j, i) | j <- [0 .. ntheta-1], i <- [0 .. nphi-1]] 
    147                 chooseRandomRay _ = do 
    148                         r1 <- drand48 
    149                         r2 <- drand48 
    150                         return $ hemisphereDirectionRay p basis r1 r2 
    151                 isHitOccluder ray = 
    152                         not $ null $ catMaybes $ map (findRayObjIntersect ray) objs 
    153                 calcOccludedRatio occludedSampleNum = 
    154                         (ntheta * nphi - toEnum occludedSampleNum) / (ntheta * nphi) 
    155                 white = Vec { x = 1, y = 1, z = 1 } 
    156  
    157 screenRay nsubsamples w h x y u v = 
    158         Ray { ray_org = zerovec, ray_dir = vnormalize (Vec { x = px, y = py, z = -1 }) } 
    159         where 
    160                 px = (x + (u / nsubsamples) - (w / 2.0)) / (w / 2.0) 
    161                 py = -(y + (v / nsubsamples) - (h / 2.0)) / (h / 2.0) 
    162  
     173  randomRays <- sequence $ replicate nsample $ randomHemisphereRay p n 
     174  let occludedNum = length [ray | ray <- randomRays, isHitOccluder objs ray] 
     175  return $ vscale white $ 1.0 - occludedRatio occludedNum 
     176    where 
     177      eps = 0.0001 
     178      n = isect_n isect 
     179      p = isect_p isect `vadd` (n `vscale` eps) 
     180      nsample = sq nao_samples 
     181      occludedRatio occludedSampleNum = 
     182          i2d occludedSampleNum / i2d nsample 
     183 
     184screenRay :: Int -> Int -> Point -> Ray 
     185screenRay w h (x, y) = 
     186    Ray { ray_org = zerovec, ray_dir = vnormalize (Vec { x = px, y = py, z = -1 }) } 
     187    where 
     188      px = (i2d x - w_2) / w_2 
     189      py = -(i2d y - h_2) / h_2 
     190      w_2 = i2d w / 2.0 
     191      h_2 = i2d h / 2.0 
     192 
     193findNearestIsect :: [Object] -> Ray -> Maybe Isect 
    163194findNearestIsect objs ray = 
    164         listToMaybe $ sort $ catMaybes $ map (findRayObjIntersect ray) objs 
     195    listToMaybe $ sort $ catMaybes $ map (findRayObjIntersect ray) objs 
    165196 
    166197data Scene = Scene { 
    167         scene_objs :: [Object] 
    168         } 
    169  
    170 raytrace shading objs ray = do 
    171         shading $ findNearestIsect objs ray 
    172  
    173 oversampling trace makeRay nsubsamples = do 
    174         mapM trace subrays >>= return . averageColor 
    175         where 
    176                 subpixels = [(u, v) | v <- [0 .. nsubsamples-1], u <- [0 .. nsubsamples-1]] 
    177                 subrays = map (uncurry makeRay) subpixels 
    178                 averageColor = flip vscale (1.0 / sq nsubsamples) . foldl vadd zerovec 
    179  
    180 calcPixelColor trace w h nsubsamples x y = do 
    181         oversampling trace makeRay nsubsamples 
    182         where 
    183                 makeRay = screenRay nsubsamples w h x y 
    184  
    185 render w h renderPixel = do 
    186         fimg <- mapM (uncurry renderPixel) screenPixels 
    187         let img = map fpcol2icol fimg 
    188         return img 
    189         where 
    190                 screenPixels = [(x, y) | y <- [0 .. h-1], x <- [0 .. w-1]] 
    191                 fpcol2icol vec = (f2i (x vec), f2i (y vec), f2i (z vec)) 
    192                 f2i = clamp 0 255 . floor . (* 255.99) 
    193  
    194 rendering scene w h nsubsamples nao_samples = do 
    195         render w h $ calcPixelColor (raytrace shading objs) w h nsubsamples 
    196         where 
    197                 shading Nothing = return zerovec 
    198                 shading (Just isect) = ambient_occlusion objs nao_samples isect 
    199                 objs = scene_objs scene 
    200  
     198      scene_objs :: [Object] 
     199    } 
     200 
     201raytrace :: (Maybe Isect -> IO FColor) -> [Object] -> Ray -> IO FColor 
     202raytrace fShading objs ray = 
     203    fShading $ findNearestIsect objs ray 
     204 
     205oversampling :: (Point -> IO FColor) -> Int -> IO FColor 
     206oversampling fCalcCol nsubsamples = 
     207  mapM fCalcCol subpixels >>= return . vaverage 
     208    where 
     209      subpixels = [(u, v) | v <- [0 .. nsubsamples-1], u <- [0 .. nsubsamples-1]] 
     210 
     211renderSubsampledRayTracePixel :: (Ray -> IO FColor) -> Int -> Int -> Int -> Point -> IO FColor 
     212renderSubsampledRayTracePixel fTrace w h nsubsamples (x, y) = 
     213  oversampling calcCol nsubsamples 
     214    where 
     215      calcCol = fTrace . screenRay subedWidth subedHeight . subedPoint 
     216      subedPoint (u, v) = (x * nsubsamples + u, y * nsubsamples + v) 
     217      subedWidth = w * nsubsamples 
     218      subedHeight = h * nsubsamples 
     219 
     220renderScreen :: Int -> Int -> (Point -> IO col) -> IO [col] 
     221renderScreen w h fRenderPixel = 
     222  mapM fRenderPixel screenPixels 
     223    where 
     224      screenPixels = [(x, y) | y <- [0 .. h-1], x <- [0 .. w-1]] 
     225 
     226renderAmbientOcclusionScene :: Scene -> Int -> Int -> Int -> Int -> IO [(Int, Int, Int)] 
     227renderAmbientOcclusionScene scene w h nsubsamples nao_samples = do 
     228  fimg <- renderScreen w h calcPixelColor 
     229  return $ map fpcol2icol fimg 
     230    where 
     231      calcPixelColor = renderSubsampledRayTracePixel (raytrace shading objs) w h nsubsamples 
     232      shading Nothing = return zerovec 
     233      shading (Just isect) = ambient_occlusion objs nao_samples isect 
     234      objs = scene_objs scene 
     235 
     236saveppm :: String -> Int -> Int -> [(Int, Int, Int)] -> IO () 
    201237saveppm fname w h img = 
    202         withBinaryFile fname WriteMode $ \fp -> do 
    203                 hPutStrLn fp "P6" 
    204                 hPutStrLn fp $ show (floor w) ++ " " ++ show (floor h) 
    205                 hPutStrLn fp "255" 
    206                 B.hPut fp $ B.pack $ map toEnum $ concatMap (\(r,g,b) -> [r,g,b]) img 
    207  
    208 init_scene = 
    209         Scene { 
    210                 scene_objs = [ 
    211                         ObjSphere (Sphere { sphere_center = Vec { x = -2.0, y = 0.0, z = -3.5 }, sphere_radius = 0.5 }), 
    212                         ObjSphere (Sphere { sphere_center = Vec { x = -0.5, y = 0.0, z = -3.0 }, sphere_radius = 0.5 }), 
    213                         ObjSphere (Sphere { sphere_center = Vec { x =  1.0, y = 0.0, z = -2.2 }, sphere_radius = 0.5 }), 
    214                         ObjPlane  (Plane  { plane_p = Vec { x = 0.0, y = -0.5, z = 0.0 }, plane_n = Vec { x = 0.0, y = 1.0, z = 0.0 } }) 
    215                         ] 
    216                 } 
     238    withBinaryFile fname WriteMode $ \fp -> do 
     239      hPutStrLn fp "P6" 
     240      hPutStrLn fp $ show w ++ " " ++ show h 
     241      hPutStrLn fp "255" 
     242      B.hPut fp $ B.pack $ map toEnum $ concatMap (\(r,g,b) -> [r,g,b]) img 
     243 
     244init_scene = Scene { 
     245               scene_objs = [ 
     246                ObjSphere (Sphere { sphere_center = Vec { x = -2.0, y = 0.0, z = -3.5 }, sphere_radius = 0.5 }), 
     247                ObjSphere (Sphere { sphere_center = Vec { x = -0.5, y = 0.0, z = -3.0 }, sphere_radius = 0.5 }), 
     248                ObjSphere (Sphere { sphere_center = Vec { x =  1.0, y = 0.0, z = -2.2 }, sphere_radius = 0.5 }), 
     249                ObjPlane  (Plane  { plane_p = Vec { x = 0.0, y = -0.5, z = 0.0 }, plane_n = Vec { x = 0.0, y = 1.0, z = 0.0 } }) 
     250               ] 
     251             } 
    217252 
    218253getParam idx defaultVal args 
    219         | length args > idx     = args !! idx 
    220         | otherwise                     = defaultVal 
     254    | length args > idx  = args !! idx 
     255    | otherwise          = defaultVal 
    221256 
    222257main = do 
    223         args <- getArgs 
    224         let width  = read $ getParam 0 "256" args 
    225         let height = read $ getParam 1 "256" args 
    226         let nsubsamples = read $ getParam 2 "2" args 
    227         let nao_samples = read $ getParam 3 "8" args 
    228  
    229         img <- rendering init_scene width height nsubsamples nao_samples 
    230         saveppm "ao.ppm" width height img 
     258  args <- getArgs 
     259  let width  = read $ getParam 0 "256" args 
     260  let height = read $ getParam 1 "256" args 
     261  let nsubsamples = read $ getParam 2 "2" args 
     262  let nao_samples = read $ getParam 3 "8" args 
     263 
     264  img <- renderAmbientOcclusionScene init_scene width height nsubsamples nao_samples 
     265  saveppm "ao.ppm" width height img