petzval/lib/Petzval/Trace.hs

-- | Utilities for full-precision raytracing

-- {-# OPTIONS_HADDOCK ignore-exports #-}
{-# LANGUAGE FlexibleContexts #-}
module Petzval.Trace
  ( Ray(..)
  , _dir, _pos
  , createRay
  , HitRecord(..)
  , TraceError(..)
  , raytrace
  , raytrace1
  -- * Ray patterns
  , hexapolarPattern
  , spiralPattern
  ) where

import Linear
import Petzval.System
import Petzval.Optics
import Numeric.AD.Mode (Scalar, Mode, auto)
import Control.Lens
import Control.Monad.State
import Control.Monad.Except
import Control.Monad.Writer

-- | 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

orError :: (MonadError e m) => Maybe a -> e -> m a
orError  = maybe throwError  (const . return)

-- | 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, in degrees
          -> V2 a    -- ^ Normalized pupil coordinates (in the range \([-1,1]\))
          -> Ray a
createRay (Just objectPlane) Pupil{position=pz,radius=pr} h (V2 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 (V2 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, relative to the vertex of the last element.
--
-- This is equivalent to `foldM raytrace1 ray system`, given an appropriate monad stack.
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 =
  runIdentity
  . flip evalStateT (1 :: Double)
  . runWriterT
  . runExceptT
  $ foldM raytrace1 ray system
  
-- | Trace a ray through a single element. Given an appropriate monad, this is a far more powerful interface to tracing than `raytrace`
raytrace1 :: ( Floating a, Ord a, Mode a, Scalar a ~ Double, Epsilon a
             , MonadState Double m
             , MonadWriter [HitRecord a] m -- ^ Tracing yields a list of 
             , MonadError TraceError m) => -- ^ This can fail
             Ray a -> Element BakedIOR a -> m (Ray a)
raytrace1 ray element = do
      n1 <- get
      let stopP = isStop element
      (nray, mnorm) <- hitTest element ray `orError` (if stopP then ElementMissed else HitStop)
      let mat@(BakedIOR _ n2)  = maybe (BakedIOR n1 n1) id $ element ^? material
      nray' <- maybe (return nray) (\normal -> refract mat normal nray `orError` TIR) mnorm
      let opl = distance (ray ^. _pos) (nray ^. _pos) * auto n1
      put n2
      tell [HitRecord { pos=(nray' ^. _pos), opl}]
      return $ nray' &_pos._z -~ element ^. thickness

-- | Spiral pattern. This is somewhat more irregular than the hexapolar pattern. The argument is the number of points
spiralPattern :: Floating a => Int -> [V2 a]
spiralPattern n = map (point . fromIntegral) [0..n-1]
  where npoints = fromIntegral n - 1
        point n = let r = sqrt (n / npoints)
                      theta = 2.3999632297286531 * n
                  in r *^ V2 (sin theta) (cos theta)
-- | A hexapolar pattern. This is rather optimally distributed; the argument is the number of rings
hexapolarPattern :: Floating a => Integer -> [V2 a]
hexapolarPattern nRings = (V2 0 0) : ( mconcat . map ring $ [0..nRings-1])
  where ring n = let point t = r *^ V2 (sin t) (cos t)
                     r = (fromIntegral n) / (fromIntegral nRings)
                 in map (point . (\x -> x / (fromIntegral) n * 3 * pi) . fromIntegral) [0 .. n-1]