crocodile-0.1.2: app/src/RayTrace.hs
-- The module where all the tracing actually happens
{-# LANGUAGE BangPatterns #-}
module RayTrace (rayTraceImage, pathTraceImage, findNearestIntersection, findAnyIntersection, GlobalIlluminationFunc) where
import PolymorphicNum
import Vector
import {-# SOURCE #-} Light
import Primitive
import Colour
import Ray
import Material
import Matrix
import Misc
import Camera
import Distribution
import SceneGraph
import Control.Parallel.Strategies
import {-# SOURCE #-} PhotonMap (PhotonMap, irradiance)
import IrradianceCache
import Control.Monad.State
import RenderContext
import System.Random.Mersenne.Pure64
import Data.List
import RussianRoulette
-- Intersect a ray against a sphere tree
intersectSphereTree :: [SphereTreeNode] -> Ray -> Maybe (Object, Double, Int) -> Maybe (Object, Double, Int)
intersectSphereTree !(node:nodes) !ray !currentHit = seq result (intersectSphereTree (newNodeList ++ nodes) newRay thisResult)
where
-- Intersect the ray with the bounding volume of this node
!result = case children node of
-- If the node has no children, don't bother with it's bounding volume and just check the object (if it has one)
[] -> case object node of
Nothing -> error "A node with no children should hold an object"
Just obj -> case primitiveClosestIntersect (primitive obj) ray obj of
-- Didn't hit the object. Retain the current hit, and continue with remaining nodes on the list
Nothing -> (currentHit, [], ray)
-- We did hit this object. Update the intersection, and continue with remaining nodes on the list
Just (objHitDistance, objHitId) -> (Just (obj, objHitDistance, objHitId), [], shortenRay ray objHitDistance)
-- We have children. In this case it makes sense to test our bounding volume
!nodeChildren -> case sphereIntersect (boundingRadius node) (boundingCentre node) ray of -- (make a sphere centred at the object's transform matrix with given radius)
-- If we do not find an intersection, we do not update the results and we offer no further nodes to be traversed, thus skipping this subtree
Nothing -> (currentHit, [], ray)
-- If we do find an intersection against the bounding volume, then we try again against the actual object (if present)
Just _ -> case object node of
Nothing -> (currentHit, nodeChildren, ray) -- No object; just pass to the children
Just obj -> case primitiveClosestIntersect (primitive obj) ray obj of
-- Didn't hit the object. Retain the current hit, but offer up the children of the node as we hit the bounding volume
Nothing -> (currentHit, nodeChildren, ray)
-- We did hit this object. Update the intersection, and continue with the bounding volume's children
Just (objHitDistance, objHitId) -> (Just (obj, objHitDistance, objHitId), nodeChildren, shortenRay ray objHitDistance)
(!thisResult, !newNodeList, !newRay) = result
intersectSphereTree [] _ !currentHit = currentHit
-- Intersect with the list of infinite objects
intersectObjectList :: [Object] -> Ray -> Maybe (Object, Double, Int) -> Maybe (Object, Double, Int)
intersectObjectList !(obj:objs) !ray !currentHit = intersectObjectList objs newRay thisResult
where
(!thisResult, !newRay) = case primitiveClosestIntersect (primitive obj) ray obj of
Nothing -> (currentHit, ray)
Just (objHitDistance, objHitId) -> (Just (obj, objHitDistance, objHitId), shortenRay ray objHitDistance)
intersectObjectList [] _ !currentHit = currentHit
-- Find the nearest intersection along a line
findNearestIntersection :: SceneGraph -> Ray -> Maybe (Object, Double, Int)
findNearestIntersection sceneGraph' !ray = case intersectObjectList (infiniteObjects sceneGraph') ray Nothing of
Just (obj, dist, objId) -> intersectSphereTree [root sceneGraph'] (shortenRay ray dist) (Just (obj, dist, objId))
Nothing -> intersectSphereTree [root sceneGraph'] ray Nothing
-- Intersect a ray against a scene graph. Return first (ie, any) hit without finding the closest
findAnyIntersectionSphereTree :: [SphereTreeNode] -> Ray -> Maybe (Object, Double)
findAnyIntersectionSphereTree (node:nodes) !ray = let sphereIntersectionResult = sphereIntersect (boundingRadius node) (boundingCentre node) ray
in case sphereIntersectionResult of -- (make a sphere centred at the object's transform matrix with given radius)
-- If we do not find an intersection, traverse to the rest of the list
Nothing -> findAnyIntersectionSphereTree nodes ray
-- If we do find an intersection against the bounding volume, then we try again against the actual object
Just _ -> case object node of
Nothing -> findAnyIntersectionSphereTree (nodes ++ children node) ray -- No object here - just offer up the children
Just obj -> case primitiveAnyIntersect (primitive obj) ray obj of
-- Didn't hit the object. Offer up the children of the scene graph node to continue with (as we did actually hit the bounding volume)
Nothing -> findAnyIntersectionSphereTree (nodes ++ children node) ray
-- We did hit this object. Update the intersection, and continue with the bounding volume's children
Just (objHitDistance, _) -> Just (obj, objHitDistance)
findAnyIntersectionSphereTree [] _ = Nothing
-- Find any intersection against an object list
findAnyIntersectionObjectList :: [Object] -> Ray -> Maybe (Object, Double)
findAnyIntersectionObjectList (obj:objs) !ray = case primitiveAnyIntersect (primitive obj) ray obj of
-- Didn't hit he object. Retain the current hit, but offer up the children of the node as we hit the bounding volume
Nothing -> findAnyIntersectionObjectList objs ray
-- We did hit this object. Update the intersection, and continue with the bounding volume's children
Just (objHitDistance, _) -> Just (obj, objHitDistance)
findAnyIntersectionObjectList [] _ = Nothing
findAnyIntersection :: SceneGraph -> Ray -> Maybe (Object, Double)
findAnyIntersection sceneGraph' !ray = case findAnyIntersectionObjectList (infiniteObjects sceneGraph') ray of
Nothing -> findAnyIntersectionSphereTree [root sceneGraph'] ray
Just x -> Just x
-- Default background colour to return when we can't match anything
defaultColour :: Direction -> Colour
defaultColour _ = colBlue
-- Accumulate the contributions of the lights
lightSurface :: [Light] -> Colour -> RenderContext -> SurfaceLocation -> Material -> Vector -> Colour
lightSurface (x:xs) !acc renderContext !posTanSpace !objMaterial !viewDirection
= let result = acc <+> emissive <+> applyLight (sceneGraph renderContext) posTanSpace objMaterial viewDirection x
emissive = emission objMaterial
in seq result (lightSurface xs result renderContext posTanSpace objMaterial viewDirection)
lightSurface [] !acc _ _ _ _ = acc
-- Abstraction to permit different GI calculations
type GlobalIlluminationFunc = (SurfaceLocation -> IrradianceCache -> Object -> RenderContext -> (Colour, IrradianceCache))
-- Photon map specific GI calculator
photonMapGlobalIllumination :: Maybe PhotonMap -> SurfaceLocation -> IrradianceCache -> Object -> RenderContext -> (Colour, IrradianceCache)
photonMapGlobalIllumination (Just photonMap) !surfaceLocation irrCache obj renderContext =
case renderMode renderContext of
PhotonMapper -> if useIrradianceCache renderContext
then query irrCache surfaceLocation irradiance'
else (fst $ irradiance photonMap (photonMapContext renderContext) (material obj) surfaceLocation, irrCache)
_ -> undefined -- Shouldn't hit this path...
where
irradiance' x = irradiance photonMap (photonMapContext renderContext) (material obj) x
photonMapGlobalIllumination _ _ irrCache _ _ = (colBlack, irrCache)
-- Null GI
nullGI :: SurfaceLocation -> IrradianceCache -> Object -> RenderContext -> (Colour, IrradianceCache)
nullGI _ irrCache _ _ = (colBlack, irrCache)
-- Retrieve the appropriate GI function and calculate the GI at a point
calculateGI :: RenderContext -> Maybe PhotonMap -> GlobalIlluminationFunc
calculateGI renderContext photonMap = case renderMode renderContext of
PhotonMapper -> photonMapGlobalIllumination photonMap
_ -> nullGI
-- Perform a full trace of a ray
type RayTraceState = State IrradianceCache Colour
traceRay :: RenderContext -> Maybe PhotonMap -> Ray -> Int -> Direction -> Double -> Double -> RayTraceState
-- Special case for lowest level of recursion (theoretically this should not get hit)
traceRay _ _ _ 0 _ _ _ = error "Should not hit this codepath"
-- Special case for penultimate level - we're not allowed to spawn rays here
traceRay renderContext photonMap !ray 1 !viewDir _ _ =
case findNearestIntersection (sceneGraph renderContext) ray of
Nothing -> return $! defaultColour (direction ray)
Just (obj, intersectionDistance, hitId) -> do
irrCache <- get
let !intersectionPoint = pointAlongRay ray intersectionDistance
let !tanSpace = primitiveTangentSpace (primitive obj) hitId intersectionPoint obj
let !(!surfaceIrradiance, !newIrrCache) = calculateGI renderContext photonMap (intersectionPoint, tanSpace) irrCache obj renderContext
-- TODO - Need to plug irradiance values into shader model correctly
let resultColour = lightSurface (lights renderContext) surfaceIrradiance renderContext (intersectionPoint, tanSpace) (material obj) viewDir
put newIrrCache
return $! resultColour
-- General case
traceRay renderContext photonMap !ray !limit !viewDir !currentIOR !accumulatedReflectivity =
case findNearestIntersection (sceneGraph renderContext) ray of
Nothing -> return $! defaultColour (direction ray)
Just (obj, intersectionDistance, hitId) -> do
-- Evaluate surface-location specific things such as shader results
let !intersectionPoint = pointAlongRay ray intersectionDistance
let !tanSpace = primitiveTangentSpace (primitive obj) hitId intersectionPoint obj
let !normal = thr tanSpace
let !incoming = Vector.negate $ direction ray
-- TODO - Need to plug irradiance values into shader model correctly
-- Evaluate result from irradiance cache
irrCache <- get
let !(!surfaceIrradiance, !irrCache') = calculateGI renderContext photonMap (intersectionPoint, tanSpace) irrCache obj renderContext
put irrCache'
let !surfaceShading = lightSurface (lights renderContext) surfaceIrradiance renderContext (intersectionPoint, tanSpace) (material obj) viewDir
-- Reflection specific code
let offsetToExterior = madd intersectionPoint normal surfaceEpsilon
let reflectionDir = normalise $ reflect incoming normal
let !shine = reflectivity $ material obj
let reflectRay = rayWithDirection offsetToExterior reflectionDir (reflectionRayLength renderContext)
let (reflection, irrCache'') = if shine > 0 && (accumulatedReflectivity * shine) > 0.03
then runState (traceRay renderContext photonMap reflectRay (limit - 1) viewDir currentIOR (accumulatedReflectivity * shine)) irrCache'
else (colBlack, irrCache')
put irrCache''
-- Refraction specific
let !eta = if enteringObject incoming normal
then currentIOR / indexOfRefraction (material obj)
else indexOfRefraction (material obj) / currentIOR
let refractionDir = normalise $ refract incoming normal eta
let offsetToInterior = madd intersectionPoint refractionDir surfaceEpsilon
let !transmittance = transmit $ material obj
let refractRay = rayWithDirection offsetToInterior refractionDir (refractionRayLength renderContext)
let (refraction, irrCache''') = if transmittance > 0
then runState (traceRay renderContext photonMap refractRay (limit - 1) viewDir (indexOfRefraction $ material obj) accumulatedReflectivity) irrCache''
else (colBlack, irrCache'')
put irrCache'''
-- Final colour combine
return $! (surfaceShading <+> reflection <*> shine <+> refraction <*> transmittance)
where
enteringObject !incoming !normal = incoming `dot3` normal > 0
-- Trace a list of distributed samples with tail recursion
rayTracePixelSample :: RenderContext -> Colour -> [Position] -> Maybe PhotonMap -> (Position, Direction) -> Double -> RayTraceState
rayTracePixelSample renderContext !acc (x:xs) photonMap !eyeViewDir !sampleWeighting =
do
irrCache <- get
let dofFocalDistance = depthOfFieldFocalDistance renderContext
let jitteredRayPosition jitter = fst eyeViewDir <+> jitter
let jitteredRayDirection jitter = normalise $ madd jitter (snd eyeViewDir) dofFocalDistance
let (sampleColour, irrCache') = runState (traceRay renderContext photonMap (rayWithDirection (jitteredRayPosition x) (jitteredRayDirection x) 100000.0) (maximumRayDepth renderContext) (snd eyeViewDir) 1 1) irrCache
let result = sampleColour <*> sampleWeighting <+> acc
put irrCache'
let (col, irrCache'') = runState (rayTracePixelSample renderContext result xs photonMap eyeViewDir sampleWeighting) irrCache'
put irrCache''
return $! col
rayTracePixelSample _ !acc [] _ _ _ = return $! acc
-- This traces for a given pixel (x, y)
rayTracePixel :: RenderContext -> Position -> Maybe PhotonMap -> Direction -> RayTraceState
rayTracePixel renderContext eye photonMap viewDirection = do
irrCache <- get
let distributedPositions = generatePointsOnSphere (numDistribSamples renderContext) (rayOriginDistribution renderContext) 12345 -- TODO Fix this magic number!
let (pixelColour, irrCache') = runState (rayTracePixelSample renderContext colBlack distributedPositions photonMap (eye, viewDirection) (1.0 / (fromIntegral . numDistribSamples $ renderContext))) irrCache
put irrCache'
return $! pixelColour
-- This function converts a pixel co-ordinate to a direction of the ray
makeRayDirection :: Int -> Int -> Camera -> (Double, Double) -> Vector
makeRayDirection !renderWidth !renderHeight camera (x, y) =
let !x' = (x / fromIntegral renderWidth) * 2.0 - 1.0
!y' = (y / fromIntegral renderHeight) * 2.0 - 1.0
!fov = 0.5 * fieldOfView camera
!fovX = tan (degreesToRadians fov)
!fovY = -tan (degreesToRadians fov)
!aspectRatio = fromIntegral renderWidth / fromIntegral renderHeight
!dirX = fovX * x'
!dirY = fovY * (-y') / aspectRatio
!rayDir = normalise (Vector dirX dirY 1 0)
in normalise $ transformVector (worldToCamera camera) rayDir
-- Generate a list of colours which contains a raytraced image. In parallel
rayTraceImage :: RenderContext -> Camera -> Int -> Int -> Maybe PhotonMap -> [Colour]
rayTraceImage renderContext camera renderWidth renderHeight photonMap = tracePixelPassingState rayDirections irrCache `using` parListChunk 256 rdeepseq
where rayDirections = [makeRayDirection renderWidth renderHeight camera (fromIntegral x, fromIntegral y) | y <- [0..(renderHeight - 1)], x <- [0..(renderWidth - 1)]]
eyePosition = Camera.position camera
irrCache = initialiseCache (sceneGraph renderContext)
-- This function is the equivalent to map, but it passes the ending state of one invocation to the next invocation
-- I'm using this rather than my mapWithState routine because I'm concerned to do so may break the parallelism of parListChunk
tracePixelPassingState !(x:xs) !st = result : tracePixelPassingState xs st'
where
(!result, !st') = runState (rayTracePixel renderContext eyePosition photonMap x) st
tracePixelPassingState [] _ = []
type PathTraceState = State PureMT Colour
pathTrace :: RenderContext -> Ray -> Int -> Direction -> Double -> Colour -> PathTraceState
-- General case
pathTrace renderContext !ray depth !viewDir !currentIOR !weight =
case findNearestIntersection (sceneGraph renderContext) ray of
Nothing -> return $! colBlack
Just (obj, intersectionDistance, hitId) ->
-- Evaluate surface-location specific things such as shader results
let intersectionPoint = pointAlongRay ray intersectionDistance
tanSpace = primitiveTangentSpace (primitive obj) hitId intersectionPoint obj
normal = thr tanSpace
incoming = Vector.negate $ direction ray
-- Thunk for emitted light
emittedLight = (emission . material) obj
-- Compute radiance at this point
radiance = lightSurface (lights renderContext) colBlack renderContext (intersectionPoint, tanSpace) (material obj) viewDir
(diffuseP, specularP) = russianRouletteCoefficients (material obj)
in do
-- TODO - Need to evaluate the shader model here!
-- Randomly decide the fate at this intersection
gen <- get
let (p, gen') = randomDouble gen
put gen'
let interaction | p < diffuseP = DiffuseReflect
| p < (diffuseP + specularP) = SpecularReflect
| otherwise = Absorb
-- Thunk for reflected light
let (randomDir, gen'') = generatePointOnHemisphere gen' 1
put gen''
-- Set up expression for reflected light
let reflectedDir
| interaction == DiffuseReflect = transformDir randomDir tanSpace
| otherwise = incoming `reflect` normal
let ray' = rayWithDirection intersectionPoint reflectedDir (rayLength ray)
let weight' = (diffuse . material) obj <*> weight <*> (normal `sdot3` reflectedDir)
let (tracedPathColour, gen''') = runState (pathTrace renderContext ray' (depth + 1) viewDir currentIOR weight') gen''
let reflectedLight = tracedPathColour <*> weight'
put gen'''
-- Have to divide by probability to correctly account for that relative proportion of the domain
return $! case interaction of DiffuseReflect -> (emittedLight <+> radiance <+> reflectedLight) </> diffuseP
SpecularReflect -> (emittedLight <+> radiance <+> reflectedLight) </> (diffuseP + specularP)
Absorb -> (emittedLight <+> radiance) </> (1 - diffuseP - specularP)
-- Path-trace a sub-sample
pathTracePixelSample :: RenderContext -> Camera -> (Int, Int) -> (Int, Int) -> (Double, Double) -> (Double, Double) -> PathTraceState
pathTracePixelSample renderContext camera xy (width, height) jitterUV stratUV = pathTrace renderContext ray 0 rayDirection 1 colWhite
where
jitteredX = (fromIntegral . fst) xy + fst stratUV + fst jitterUV
jitteredY = (fromIntegral . snd) xy + snd stratUV + snd jitterUV
rayDirection = makeRayDirection width height camera (jitteredX, jitteredY)
ray = rayWithDirection (Camera.position camera) rayDirection (farClip camera)
-- Path-trace a pixel. Do stratified sub-sampling
pathTracePixel :: RenderContext -> Camera -> (Int, Int) -> (Int, Int) -> PathTraceState
pathTracePixel renderContext camera pixelCoords renderTargetSize =
do
-- Work out the jittered UV offsets
gen <- get
let (offsetUVs, gen') = runState (generateRandomUVs numPathTraceSamples) gen
let offsetUVs' = map (\(u, v) -> (u * du, v * dv)) offsetUVs
put gen'
-- Mung it all together
let (pixelSamples, gen'') = zipWithState (pathTracePixelSample renderContext camera pixelCoords renderTargetSize) offsetUVs' stratifiedCentres gen'
put gen''
return $! foldl' (\x y -> x <*> weight <+> y) colBlack pixelSamples
where
-- Total number of samples to take
numPathTraceSamplesRoot = 6 :: Int
numPathTraceSamples = numPathTraceSamplesRoot * numPathTraceSamplesRoot
weight = (1.0 :: Double) / fromIntegral numPathTraceSamples
-- Work out a set of stratified centres to jitter from
du = (1.0 :: Double) / fromIntegral numPathTraceSamplesRoot
dv = du
stratifiedCentres = [(fromIntegral x * du, fromIntegral y * dv) | y <- [0..(numPathTraceSamplesRoot - 1)], x <- [0..(numPathTraceSamplesRoot - 1)]]
pathTraceImage :: RenderContext -> Camera -> Int -> Int -> [Colour]
pathTraceImage renderContext camera renderWidth renderHeight = zipWith
(\x y -> evalState (pathTracePixel renderContext camera x (renderWidth, renderHeight)) (pureMT y))
[(x, y) | y <- [0..(renderHeight - 1)], x <- [0..(renderWidth - 1)]]
[1..]
`using` parListChunk 256 rdeepseq
-- TODO Re-unify the path tracer and ray tracer/photon map code paths as much as practical to ease maintenance (?)