127 lines
5.1 KiB
Haskell
127 lines
5.1 KiB
Haskell
-- | Utilities for full-precision raytracing
|
|
|
|
-- {-# OPTIONS_HADDOCK ignore-exports #-}
|
|
module Petzval.Trace
|
|
( Ray(..)
|
|
, _dir, _pos
|
|
, createRay
|
|
, HitRecord(..)
|
|
, TraceError(..)
|
|
, raytrace
|
|
) where
|
|
|
|
import Linear
|
|
import Petzval.System
|
|
import Petzval.Optics
|
|
import Numeric.AD.Mode (Scalar, Mode, auto)
|
|
import Control.Lens
|
|
|
|
-- | A ray. The first argument is the direction, and the second
|
|
data Ray a = Ray (V3 a) (V3 a)
|
|
deriving (Show, Eq)
|
|
|
|
_dir, _pos :: Lens' (Ray a) (V3 a)
|
|
-- | The direction of a ray
|
|
_dir = lens (\(Ray _ dir) -> dir) (\(Ray pos _) -> Ray pos)
|
|
-- | The position of a ray
|
|
_pos = lens (\(Ray pos _) -> pos) (\(Ray _ pos) dir -> Ray dir pos)
|
|
|
|
toMaybe :: Bool -> a -> Maybe a
|
|
toMaybe False = const Nothing
|
|
toMaybe True = Just
|
|
|
|
orLeft :: Maybe a -> b -> Either b a
|
|
orLeft = maybe Left (const . Right)
|
|
|
|
-- | Create a ray for a given field angle and pupil position.
|
|
--
|
|
-- * The first argument is the image plane position. If `Nothing`, the object is at infinity.
|
|
--
|
|
-- * The second argument is the entrance pupil to aim at
|
|
--
|
|
-- * The third argument is the field angle
|
|
--
|
|
-- * The fourth argument is the normalized pupil coordinates (in the range of \([-1,1]\))
|
|
createRay :: (RealFloat a, Mode a, Epsilon a)
|
|
=> Maybe a -- ^ The image plane position. If `Nothing`, the object is at infinity
|
|
-> Pupil a -- ^ The entrance pupil to aim at
|
|
-> a -- ^ Field angle
|
|
-> (a,a) -- ^ Normalized pupil coordinates (in the range \([-1,1]\))
|
|
-> Ray a
|
|
createRay (Just objectPlane) Pupil{position=pz,radius=pr} h (px, py) =
|
|
Ray source (normalize $ target ^-^ source)
|
|
where dz = pz - objectPlane
|
|
source = V3 0 (dz * tan h) objectPlane
|
|
target = V3 (px * pr) (py * pr) pz
|
|
createRay Nothing Pupil{position=pz,radius=pr} h (px, py) = Ray source (normalize $ target ^-^ source)
|
|
where h' = (pi * (-abs h) / 180) -- field angle in rad
|
|
dy = (V3 0 (cos h') (-sin h')) `project` (V3 0 (py * pr) 0)
|
|
dz = V3 0 (pz * tan h') (pz * cos h')
|
|
|
|
source = dy ^+^ dz
|
|
target = V3 (px * pr) (py * pr) pz
|
|
|
|
|
|
|
|
hitTest :: (Floating a, Ord a, Mode a, Epsilon a) => Element mat a -> Ray a -> Maybe (Ray a, Maybe (V3 a))
|
|
hitTest Stop{_outsideRadius} (Ray pos dir) =
|
|
toMaybe pass $ (Ray npos dir, Nothing)
|
|
where dz = -pos ^. _z / dir ^. _z
|
|
npos = pos ^+^ (dir ^* dz)
|
|
pass = quadrance (npos ^. _xy) < _outsideRadius ^ 2
|
|
hitTest Surface{_roc, _outsideRadius} (Ray pos dir) =
|
|
toMaybe (hit1 && hit2) (Ray npos dir, Just normal)
|
|
where origin = dir & _z -~ _roc
|
|
a = dir `dot` dir
|
|
b = (dir `dot` origin) * 2
|
|
c = (origin `dot` origin) - _roc ^ 2
|
|
det = b^2 - 4 * a * c
|
|
hit1 = det >= 0
|
|
p2 = sqrt det
|
|
sa = (p2 - b) / 2 / a
|
|
sb = (-p2 - b) / 2 / a
|
|
s1 = min sa sb
|
|
s2 = max sa sb
|
|
dist = if s1 >= 0.001 then s1 else s2
|
|
normal = normalize $ origin ^+^ dir ^* dist
|
|
npos = pos ^+^ dir ^* dist
|
|
hit2 = (quadrance $ npos ^. _xy) <= _outsideRadius^2
|
|
|
|
refract :: (Floating a, Ord a, Mode a, Scalar a ~ Double, Epsilon a) => BakedIOR -> V3 a -> Ray a -> Maybe (Ray a)
|
|
refract (BakedIOR n1 n2) normal (Ray pos incident) =
|
|
let ni = normal `dot` incident
|
|
mu = auto $ n1 / n2
|
|
det = 1 - mu^2 * (1 - ni^2)
|
|
in toMaybe (det >= 0) $ Ray pos $ mu *^ incident + (sqrt det - mu * ni) *^ normal
|
|
|
|
-- | The interaction of a ray with a particular element
|
|
data HitRecord a = HitRecord { pos :: V3 a -- ^ Position of the hit
|
|
, opl :: a -- ^ Optical path length from the last hit to here
|
|
}
|
|
deriving (Show)
|
|
|
|
-- | How a ray failed to complete a trace
|
|
data TraceError = HitStop -- ^ Ray passed outside the aperture stop
|
|
| ElementMissed -- ^ The ray missed an element completely
|
|
| TIR -- ^ The ray hit an element so obliquely that it
|
|
-- suffered from total internal reflection
|
|
deriving (Show, Eq)
|
|
|
|
-- | Trace a ray through the give system. Returns the ray after the last element (rebased relative to the beginning of the optical system)
|
|
raytrace :: (Floating a, Ord a, Mode a, Scalar a ~ Double, Epsilon a)
|
|
=> [Element BakedIOR a] -- ^ The system to trace
|
|
-> Ray a -- ^ The initial ray
|
|
-> Either TraceError (Ray a, [HitRecord a])
|
|
raytrace system ray = trace' 1 ray system
|
|
where -- trace' :: Double -> Ray a -> [Element BakedIOR a] -> Either TraceError [HitRecord a]
|
|
trace' n1 ray@(Ray pos dir) (element:elements) = do
|
|
let stopP = isStop element
|
|
(nray, mnorm) <- hitTest element ray `orLeft` (if stopP then ElementMissed else HitStop)
|
|
let mat@(BakedIOR _ n2) = maybe (BakedIOR n1 n1) id $ element ^? material
|
|
nray' <- maybe (Right nray) (\normal -> refract mat normal nray `orLeft` TIR) mnorm
|
|
|
|
let opl = distance pos (nray ^. _pos) * auto n1
|
|
(fray, rest) <- trace' n2 (nray'&_pos._z -~ element ^. thickness) elements
|
|
|
|
return (fray & _pos._z +~ element ^. thickness, HitRecord { pos=(nray' ^. _pos), opl} : rest)
|