Got real tracing, seidel calulations working

lib/Petzval/Optics.hs

@@ -11,6 +11,7 @@ module Petzval.Optics
   , ConstMat(..)
     -- * Optical elements
   , Element(..)
+  , isStop, isSurface
   , thickness
   , outsideRadius
   , material
@@ -64,6 +65,15 @@ data Element mat fp =
       }
                     deriving (Show)
 makeLenses ''Element
+
+-- | Determine if an element is a stop
+isStop :: Element mat a -> Bool
+isStop Stop{} = True
+isStop _ = False
+-- | Determine if an element is a refractive surface
+isSurface :: Element mat a -> Bool
+isSurface Surface{} = True
+isSurface _ = False
     {-
 -- | The space after the current element
 thickness :: Lens (Element mat a) (Element mat a) a a
@@ -94,12 +104,14 @@ liftFp inj (s@Surface{_thickness, _outsideRadius, _roc, _material}) = (\t' or' r
 liftFp inj (s@Stop{_thickness, _outsideRadius}) = (\t' or' -> Stop{_thickness=t', _outsideRadius=or'}) <$> inj _thickness <*> inj _outsideRadius
 {-# INLINE liftFp #-}
 
-specialize :: (Material mat) => Double -> [Element mat a] -> [Element Double a]
+-- | Specialize a system for a specific wavelength
-specialize wavelength = (each.material) %~ (iorAtWavelength wavelength)
+specialize :: (Material mat) => Double -> [Element mat a] -> [Element ConstMat a]
+specialize wavelength = (each.material) %~ (ConstMat . iorAtWavelength wavelength)
+-- | Annotate each material with the incoming index of refraction
 bake :: (Material mat) => Double -> [Element mat a] -> [Element BakedIOR a]
 bake wavelength mat =
   snd $
-  mapAccumLOf (each.material) (\n1 n2 -> (n2, BakedIOR n1 n2)) 1 $
+  mapAccumLOf (each.material) (\n1 (ConstMat n2) -> (n2, BakedIOR n1 n2)) 1 $
   specialize wavelength mat
 
 

lib/Petzval/System.hs

@@ -1,9 +1,8 @@
+-- | General utilities for dealing with full lens systems (i.e., the composition of multiple refracting surfaces)
 module Petzval.System
   ( Pupil(..)
   , entrancePupil
-  , createRay
   , seidel, Seidel(..)
-  , Ray
   ) where
 
 import Petzval.Optics
@@ -15,7 +14,6 @@ import Data.List
 
 import qualified Debug.Trace
 
-data Ray a = Ray (V3 a) (V3 a)
 
 splitSystem :: [ Element mat a ] -> ([Element mat a], Maybe (Element mat a), [Element mat a])
 splitSystem (s@Stop{}:rest) = ([], Just s, rest)
@@ -23,6 +21,7 @@ splitSystem [] = ([], Nothing, [])
 splitSystem (s:rest) = (s:pfx, stop, sfx)
   where (pfx, stop, sfx) = splitSystem rest
 
+-- | A pupil (real or virtual, entrance or exit)
 data Pupil a = Pupil { radius :: a, position :: a }
   deriving (Show)
 
@@ -33,24 +32,13 @@ pupil (stop@Stop{_outsideRadius}) subsystem =
       --msg = intercalate " " ["ptrace: ", show (rtm !* V2 (_outsideRadius / my ) 0)]
   in Pupil { radius = _outsideRadius / my
            , position = -cy / cu }
+
+-- | Compute the entrance pupil of a system viewed from infinity
 entrancePupil :: (RealFloat a, Mode a, Scalar a ~ Double, Epsilon a, Show a) => [Element BakedIOR a] -> Pupil a
 entrancePupil system = pupil stop prefix
   where (prefix, (Just stop), _) = splitSystem system
 
-createRay :: (RealFloat a, Mode a, Scalar a ~ Double, Epsilon a) => Maybe a -> Pupil a -> a -> (a,a) -> Ray a
+-- | The set of seidel coefficients of a lens system
-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
-
 data Seidel a = Seidel
                    { sphr, coma, asti, fcur, dist :: a
                    -- , c1, c2 :: a
@@ -114,9 +102,7 @@ seidel' rays (s@Surface{_material=BakedIOR n1 n2,_roc=roc}) = Debug.Trace.trace
         -- cbas = h  (n2 - n1) / n1
         -- c1 = 
         
-        
+-- | Compute the seidel coefficients of a system given the pupil and field angle
-  
-
 seidel :: (RealFloat a, Mode a, Scalar a ~ Double, Epsilon a, Show a) => Pupil a -> Double -> [Element BakedIOR a] -> [Seidel a]
 seidel pupil fieldAngle =
   snd . mapAccumL seidel' rays

lib/Petzval/Trace.hs

@@ -0,0 +1,126 @@
+-- | 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)

petzval-hs.cabal

@@ -22,7 +22,7 @@ maintainer:         thequux@thequux.com
 extra-source-files: CHANGELOG.md
 
 common base
-    default-extensions: GADTs, NamedFieldPuns
+    default-extensions: GADTs, NamedFieldPuns, ScopedTypeVariables, TemplateHaskell
     default-language: Haskell2010
                     
 executable petzval-hs
@@ -45,6 +45,7 @@ library
         Petzval.Optics
         Petzval.Optics.RTM
         Petzval.System
+        Petzval.Trace
     hs-source-dirs: lib
     build-depends:    base   ^>=4.15.1.0,
                       lens   ^>=5.2,