Lots of changes

app/Main.hs

@@ -1,12 +1,15 @@
 module Main where
 
 import Control.Lens
+import Criterion.Main
 import Data.Either
 import Petzval.Optics
 import Petzval.Optics.RTM
 import Petzval.System
 import Petzval.Trace
 import Petzval.Calculations
+import Petzval.Optimization
+import Petzval.Types
 import Linear
 import Numeric.AD.Mode (Scalar, Mode)
 
@@ -23,8 +26,8 @@ n_ssk8 = SellemeierMat [ (1.44857867, 1.17965926e-01)
 
 system1 =
   [ Stop{_thickness = 0, _outsideRadius=5}
-  , Surface{_material = bk7, _thickness = 10, _roc = 100, _outsideRadius=0}
+  , Surface{_material = bk7, _thickness = 10, _roc = 100, _outsideRadius=10}
-  , Surface{_material = air, _thickness = 95, _roc = -100, _outsideRadius=0}
+  , Surface{_material = air, _thickness = 95, _roc = -100, _outsideRadius=10}
   ]
 
 system2 =
@@ -40,7 +43,7 @@ system2 =
   where s_sk16 = ConstMat 1.620408
         s_f4 = ConstMat 1.616589
 
-merit :: (Material mat, RealFloat a, Scalar a ~ Double, Epsilon a, Mode a) => [Element mat a] -> a
+merit :: (Material mat, Calcuable a) => [Element mat a] -> a
 -- merit :: (Material mat) => [Element mat Double] -> Double
 merit system = result 
   where baked = bake 0.587618 system
@@ -53,9 +56,13 @@ merit system = result
                   . rights
                   . map fst
                   . map (raytrace baked . createRay Nothing ep angle)
-                  $ spiralPattern 100)
+                  $ spiralPattern 10)
           $ [0, 7.07, 15]
 
+type Id a = a -> a
+
+
+
 main :: IO ()
 main =
   let baked = bake 0.587618 system2
@@ -65,3 +72,10 @@ main =
     putStrLn $ show ep
     putStrLn $ show $ rearFocalLength $ lensRTM baked
     putStrLn $ show $ mconcat (seidel ep fa baked :: [Seidel Double])
+    defaultMain [
+      bgroup "merit" [ bench "system1" $ nf merit (system1 :: [Element SellemeierMat Double])
+                     , bench "system2" $ nf merit (system2 :: [Element ConstMat Double])
+                     ]
+      , bgroup "ad" [ bench "system1" $ nf (gradientAt (ix 1.thickness) merit) system1
+                    ]
+      ]

flake.nix

@@ -18,6 +18,7 @@
         devShell = pkgs.mkShell {
           buildInputs = with pkgs; [
             cabal-install
+            stack
             haskell-language-server
             ghc
           ];

lib/Petzval/Calculations.hs

@@ -0,0 +1,21 @@
+module Petzval.Calculations where
+
+import Control.Arrow
+import Data.Foldable
+import Linear
+
+
+
+foldl1' fn (x:xs) = foldl' fn x xs
+
+bimap2 :: (a -> b -> c) -> (a' -> b' -> c') -> (a,a') -> (b,b') -> (c,c')
+bimap2 af bf (a,b) (a',b') = (af a a', bf b b')
+
+rmsSize' :: Floating a => V2 a -> [V2 a] -> a
+rmsSize' centroid = sqrt . uncurry (/) . foldl1' (bimap2 (+) (+))  . map (quadrance . (^-^ centroid) &&& (const 1))
+
+rmsSize :: Floating a => [V2 a] -> a
+rmsSize = centroid >>= rmsSize'
+
+centroid :: Fractional a => [V2 a] -> V2 a
+centroid = uncurry (^/) . foldl1' (bimap2 (^+^) (+)) . map (flip (,) 1)

lib/Petzval/Examples.hs

@@ -0,0 +1,69 @@
+module Petzval.Examples where
+
+import Control.Lens
+import Data.Either
+import Petzval.Optics
+import Petzval.Optics.RTM
+import Petzval.System
+import Petzval.Trace
+import Petzval.Calculations
+import Linear
+import Numeric.AD.Mode (Scalar, Mode)
+
+import qualified Debug.Trace 
+
+bk7 = SellemeierMat [ (1.03961212,  6.00069867e-3 )
+                    , (0.231792344, 2.00179144e-2 )
+                    , (1.01046945,  103.560653 )]
+n_sk16 = SellemeierMat [ (1.343177740, 0.007046873)
+                       , (0.241144399, 0.0229005000)
+                       , (0.994317969, 92.75085260)]
+
+n_ssk8 = SellemeierMat [ (1.44857867, 1.17965926e-01)
+                       , (1.06937528, 8.69310149E-03)
+                       , (4.21566593E-02, 1.11300666E+02) ]
+
+system1 :: Num a => [Element SellemeierMat a]
+system1 =
+  [ Stop{_thickness = 0, _outsideRadius=5}
+  , Surface{_material = bk7, _thickness = 10, _roc = 100, _outsideRadius=10}
+  , Surface{_material = air, _thickness = 95, _roc = -100, _outsideRadius=10}
+  ]
+
+system2 :: Fractional a => [Element ConstMat a]
+system2 =
+  [
+    Surface{_material = s_sk16, _outsideRadius=11.5, _roc=42.98790, _thickness = 4}
+  , Surface{_material = air, _outsideRadius = 11.5, _roc = -248.07740, _thickness = 10.51018}
+  , Surface{_material = s_f4, _outsideRadius = 9.852, _roc = -38.21035, _thickness = 2.5}
+  , Surface{_material = air, _outsideRadius = 8.885, _roc = 43.95894, _thickness = 0}
+  , Stop{_outsideRadius = 8.6762522794, _thickness = 9.86946}
+  , Surface{_material=s_sk16, _outsideRadius = 11, _roc=656.66349, _thickness = 4.5}
+  , Surface{_material = air, _outsideRadius = 11, _roc = -33.50754, _thickness = 86.48643}
+  ]
+  where s_sk16 = ConstMat 1.620408
+        s_f4 = ConstMat 1.616589
+
+trace1 :: Show a => String -> a -> a
+trace1 msg = (msg++) . show >>= Debug.Trace.trace
+trace1 _ = id
+merit :: (Material mat, RealFloat a, Scalar a ~ Double, Epsilon a, Mode a, Show a) => [Element mat a] -> a
+-- merit :: (Material mat) => [Element mat Double] -> Double
+merit system = result 
+  where baked = bake 0.587618 system
+        ep = trace1 "EP: " $ entrancePupil baked
+        result = sum
+          . map (^2)
+          . map (\ angle -> rmsSize
+                  . toListOf (each._pos._xy)
+                  -- . (\x -> x :: [Ray a])
+                  . rights
+                  . map fst
+                  . map (trace1 "Trace: " . raytrace baked . trace1 "Ray: " . createRay Nothing ep angle . trace1 "Pupil: ")
+                  $ spiralPattern 10)
+          $ [0, 7.07, 15]
+
+cast1 system = result
+  where baked = bake 0.587618 system
+        ep = trace1 "EP: " $ entrancePupil baked
+        result = raytrace baked . trace1 "Ray: " . createRay Nothing ep 15 . trace1 "Pupil: " 

lib/Petzval/Optics.hs

@@ -100,8 +100,8 @@ material inj (s@Stop{_thickness, _outsideRadius}) = (const Stop{_thickness, _out
 -}
 -- | Translate a lens element from one FP type to another. E.g., can be used to convert from scalars to the types in an automatic differentiation tool.
 liftFp :: (Applicative f) => (fp -> f fp') -> Element m fp -> f (Element m fp')
-liftFp inj (s@Surface{_thickness, _outsideRadius, _roc, _material}) = (\t' or' roc' -> Surface{_thickness=t', _outsideRadius=or', _roc=roc', _material}) <$> inj _thickness <*> inj _outsideRadius <*> inj _roc
+liftFp inj (s@Stop{_thickness, _outsideRadius}) = Stop <$> inj _thickness <*> inj _outsideRadius
-liftFp inj (s@Stop{_thickness, _outsideRadius}) = (\t' or' -> Stop{_thickness=t', _outsideRadius=or'}) <$> inj _thickness <*> inj _outsideRadius
+liftFp inj (s@Surface{_thickness, _outsideRadius, _roc, _material}) = Surface <$> inj _thickness <*> inj _outsideRadius <*> inj _roc <*> pure _material
 {-# INLINE liftFp #-}
 
 -- | Specialize a system for a specific wavelength

lib/Petzval/Optimization.hs

@@ -0,0 +1,33 @@
+{-# LANGUAGE NoMonomorphismRestriction #-}
+module Petzval.Optimization where
+
+import Control.Lens
+import Control.Lens.Unsound (adjoin)
+import Numeric.AD.Mode
+import Numeric.AD.Mode.Reverse.Double
+import Petzval.Optics
+import Petzval.Types
+
+-- | A set of modifiable parts of a lens system, expressed as a traversal.
+-- The recommended way of generating a variable set is with the following construction:
+-- @
+-- vars = ix 1 . roc `adjoin` ix 2 . thickness
+-- @
+type VariableSet = forall mat a. Traversal' [Element mat a] a
+type AdMode s = ReverseDouble s
+
+extractVars :: VariableSet -> [Element mat a] -> [a]
+extractVars vars system = system ^. partsOf vars
+
+setVars :: VariableSet -> [Element mat a] -> [a] -> [Element mat a]
+setVars vars system vals = system & partsOf vars .~ vals
+
+gradientAt :: VariableSet -- ^ The set of independent variables
+           -> (forall a. Calcuable a => [Element mat a] -> a) -- ^ merit function
+           -> [Element mat Double] -- ^ The system
+           -> (Double, [Double]) -- ^ The gradient
+gradientAt vars merit system =  grad' (merit . setVars vars (system & each.liftFp %~ auto)) (extractVars vars system)
+           
+
+
+-- instances

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.
@@ -62,38 +67,42 @@ createRay (Just objectPlane) Pupil{position=pz,radius=pr} h (V2 px py) =
         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
+  where dz = pos ^. _z / dir ^. _z
-        npos = pos ^+^ (dir ^* dz)
+        npos = pos ^-^ (dir ^* dz)
-        pass = quadrance (npos ^. _xy) < _outsideRadius ^ 2
+        pass = quadrance (npos ^. _xy) <= _outsideRadius ^ 2
-hitTest Surface{_roc, _outsideRadius} (Ray pos dir) =
+hitTest Surface{_roc, _outsideRadius} ray@(Ray pos dir) =
   toMaybe (hit1 && hit2) (Ray npos dir, Just normal)
-  where origin = dir & _z -~ _roc
+  where origin = pos & _z -~ _roc
-        a = dir `dot` dir
+        !a = dir `dot` dir
-        b = (dir `dot` origin) * 2
+        !b = (dir `dot` origin) * 2
-        c = (origin `dot` origin) - _roc ^ 2
+        !c = (origin `dot` origin) - _roc ^ 2
-        det = b^2 - 4 * a * c
+        !det = b^2 - 4 * a * c
-        hit1 = det >= 0
+        !hit1 = det >= 0
-        p2 = sqrt det
+        !p2 = sqrt det
-        sa = (p2 - b) / 2 / a
+        !sa = (p2 - b) / 2 / a
-        sb = (-p2 - b) / 2 / a
+        !sb = (-p2 - b) / 2 / a
-        s1 = min sa sb
+        !s1 = min sa sb
-        s2 = max sa sb
+        !s2 = max sa sb
-        dist = if s1 >= 0.001 then s1 else s2
+        !dist = if s1 >= -0.001 then s1 else s2
-        normal = normalize $ origin ^+^ dir ^* dist
+        !normal0 = normalize $ origin ^+^ dir ^* dist
-        npos = pos ^+^ dir ^* dist
+        !normal = if (normal0 ^. _z < 0) then -normal0 else normal0
-        hit2 = (quadrance $ npos ^. _xy) <= _outsideRadius^2
+        !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, mnorm) <- hitTest element ray `orError` (if stopP then HitStop else ElementMissed)
-      nray' <- maybe (return nray) (\normal -> refract mat normal nray `orError` TIR) mnorm
+      let !mat@(BakedIOR _ n2)  = maybe (BakedIOR n1 n1) id $ element ^? material
-      let opl = distance (ray ^. _pos) (nray ^. _pos) * auto n1
+      !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

lib/Petzval/Types.hs

@@ -0,0 +1,16 @@
+{-# LANGUAGE FlexibleInstances, UndecidableInstances #-}
+module Petzval.Types(Calcuable) where
+
+import Numeric.AD.Mode
+import Numeric.AD.Mode.Reverse.Double
+import Numeric.AD.Internal.Reverse.Double
+import Data.Reflection (Reifies)
+import Linear(Epsilon(..))
+
+
+class (RealFloat n, Scalar n ~ Double, Epsilon n, Mode n) => Calcuable n
+
+instance (RealFloat n, Scalar n ~ Double, Epsilon n, Mode n) => Calcuable n
+
+instance Reifies s Tape => Epsilon (ReverseDouble s) where
+  nearZero s = abs s <= 1e-12

petzval-hs.cabal

@@ -22,7 +22,7 @@ maintainer:         thequux@thequux.com
 extra-source-files: CHANGELOG.md
 
 common base
-    default-extensions: GADTs, NamedFieldPuns, ScopedTypeVariables, TemplateHaskell
+    default-extensions: GADTs, NamedFieldPuns, ScopedTypeVariables, TemplateHaskell, RankNTypes
     default-language: Haskell2010
                     
 executable petzval-hs
@@ -37,10 +37,26 @@ executable petzval-hs
     build-depends:    base   ^>=4.15.1.0,
                       ad     ^>=4.5.2,
                       linear ^>=1.21,
-                      lens   ^>=5.2,
+                      lens   ^>=5.0,
+                      criterion ^>=1.5,
                       petzval-hs
     hs-source-dirs:   app
 
+executable petzval-prof
+    import: base
+    main-is: Main.hs
+    build-depends:    base   ^>=4.15.1.0,
+                      ad     ^>=4.5.2,
+                      linear ^>=1.21,
+                      lens   ^>=5.0,
+                      criterion ^>=1.5,
+                      petzval-hs
+    hs-source-dirs:   app
+    ghc-options:
+      -O2
+      -threaded
+      -fprof-auto
+      -rtsopts
       
 library
     import: base
@@ -50,11 +66,15 @@ library
         Petzval.System
         Petzval.Trace
         Petzval.Calculations
+        Petzval.Optimization
+        Petzval.Types
     other-modules:
         Petzval.Internal.Vec
     hs-source-dirs: lib
     build-depends:    base   ^>=4.15.1.0,
-                      lens   ^>=5.2,
+                      lens   ^>=5.0,
                       ad     ^>=4.5.2,
                       linear ^>=1.21,
-                      mtl    ^>=2.2
+                      reflection ^>=2.1,
+                      mtl    ^>=2.2,
+                      deepseq ^>=1.4