diff --git a/app/Main.hs b/app/Main.hs index 3969ca5..95c6fa4 100644 --- a/app/Main.hs +++ b/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 = air, _thickness = 95, _roc = -100, _outsideRadius=0} + , Surface{_material = bk7, _thickness = 10, _roc = 100, _outsideRadius=10} + , 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 + ] + ] diff --git a/flake.nix b/flake.nix index 4be0ecd..89d8f52 100644 --- a/flake.nix +++ b/flake.nix @@ -18,6 +18,7 @@ devShell = pkgs.mkShell { buildInputs = with pkgs; [ cabal-install + stack haskell-language-server ghc ]; diff --git a/lib/Petzval/Calculations.hs b/lib/Petzval/Calculations.hs new file mode 100644 index 0000000..748e964 --- /dev/null +++ b/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) diff --git a/lib/Petzval/Examples.hs b/lib/Petzval/Examples.hs new file mode 100644 index 0000000..bdbf8f7 --- /dev/null +++ b/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: " diff --git a/lib/Petzval/Optics.hs b/lib/Petzval/Optics.hs index 276287c..50b2298 100644 --- a/lib/Petzval/Optics.hs +++ b/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}) = (\t' or' -> Stop{_thickness=t', _outsideRadius=or'}) <$> inj _thickness <*> inj _outsideRadius +liftFp inj (s@Stop{_thickness, _outsideRadius}) = Stop <$> 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 diff --git a/lib/Petzval/Optimization.hs b/lib/Petzval/Optimization.hs new file mode 100644 index 0000000..17236db --- /dev/null +++ b/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 diff --git a/lib/Petzval/Trace.hs b/lib/Petzval/Trace.hs index c44d9c3..bd8bbfc 100644 --- a/lib/Petzval/Trace.hs +++ b/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 diff --git a/lib/Petzval/Types.hs b/lib/Petzval/Types.hs new file mode 100644 index 0000000..060f758 --- /dev/null +++ b/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 diff --git a/petzval-hs.cabal b/petzval-hs.cabal index 45aed00..77f7729 100644 --- a/petzval-hs.cabal +++ b/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,11 +37,27 @@ 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 exposed-modules: @@ -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