Lots of changes

lib/Petzval/Trace.hs

@@ -1,7 +1,6 @@
 -- | Utilities for full-precision raytracing
-
 -- {-# OPTIONS_HADDOCK ignore-exports #-}
-{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleContexts, BangPatterns, DeriveAnyClass, DeriveGeneric #-}
 module Petzval.Trace
   ( Ray(..)
   , _dir, _pos
@@ -19,14 +18,17 @@ import Linear
 import Petzval.System
 import Petzval.Optics
 import Numeric.AD.Mode (Scalar, Mode, auto)
+import Control.DeepSeq
 import Control.Lens
 import Control.Monad.State
 import Control.Monad.Except
 import Control.Monad.Writer
+import GHC.Generics
+import qualified Debug.Trace
 
 -- | A ray. The first argument is the direction, and the second 
 data Ray a = Ray (V3 a) (V3 a)
-  deriving (Show, Eq)
+  deriving (Show, Eq, Generic, NFData)
 
 _dir, _pos :: Lens' (Ray a) (V3 a)
 -- | The direction of a ray
@@ -41,6 +43,9 @@ toMaybe True = Just
 orError :: (MonadError e m) => Maybe a -> e -> m a
 orError  = maybe throwError  (const . return)
 
+forceRay :: Ray a -> Ray a
+forceRay ray@(Ray (V3 !px !py !pz) (V3 !dx !dy !dz)) = ray
+
 -- | 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.
@@ -55,45 +60,49 @@ createRay :: (RealFloat a, Mode a, Epsilon a)
           -> 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
+          -> 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
+  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')
+        dz =  V3 0 (tan h') 1
       
-        source = dy ^+^ dz
+        source = (dy ^-^ dz * 10) & _x .~ (px * pr)
         target = V3 (px * pr) (py * pr) pz
 
+trace1 :: Show a => String -> a -> a
+trace1 msg = (msg++) . show >>= Debug.Trace.trace
 
 
 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) =
+  where dz = pos ^. _z / dir ^. _z
+        npos = pos ^-^ (dir ^* dz)
+        pass = quadrance (npos ^. _xy) <= _outsideRadius ^ 2
+hitTest Surface{_roc, _outsideRadius} ray@(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
+  where origin = pos & _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
+        !normal0 = normalize $ origin ^+^ dir ^* dist
+        !normal = if (normal0 ^. _z < 0) then -normal0 else normal0
+        !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) =
@@ -103,7 +112,7 @@ refract (BakedIOR n1 n2) normal (Ray pos incident) =
   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
+data HitRecord a = HitRecord { pos :: Ray a -- ^ Position of the hit
                              , opl :: a    -- ^ Optical path length from the last hit to here
                              }
   deriving (Show)
@@ -138,12 +147,13 @@ raytrace1 :: ( Floating a, Ord a, Mode a, Scalar a ~ Double, Epsilon 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
+      
+      (nray, mnorm) <- hitTest element ray `orError` (if stopP then HitStop else ElementMissed)
+      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}]
+      tell [HitRecord { pos=nray', 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