Reworked raytracing to expose single-surface trace

lib/Petzval/Trace.hs

@@ -1,6 +1,7 @@
 -- | Utilities for full-precision raytracing
 
 -- {-# OPTIONS_HADDOCK ignore-exports #-}
+{-# LANGUAGE FlexibleContexts #-}
 module Petzval.Trace
   ( Ray(..)
   , _dir, _pos
@@ -8,6 +9,7 @@ module Petzval.Trace
   , HitRecord(..)
   , TraceError(..)
   , raytrace
+  , raytrace1
   -- * Ray patterns
   , hexapolarPattern
   , spiralPattern
@@ -18,6 +20,9 @@ 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)
@@ -33,8 +38,8 @@ toMaybe :: Bool -> a -> Maybe a
 toMaybe False = const Nothing
 toMaybe True = Just
 
-orLeft :: Maybe a -> b -> Either b a
-orLeft  = maybe Left  (const . Right)
+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.
 --
@@ -48,7 +53,7 @@ orLeft  = maybe Left  (const . Right)
 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       -- ^ 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) = 
@@ -110,23 +115,36 @@ data TraceError = HitStop       -- ^ Ray passed outside the aperture stop
                                 --   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)
+-- | 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 = 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
+         -> (Either TraceError (Ray a), [HitRecord a])
+raytrace system ray =
+  runIdentity
+  . flip evalStateT (1 :: Double)
+  . runWriterT
+  . runExceptT
+  $ foldM raytrace1 ray system
   
-          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)
+-- | 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]