All compile errors fixed, but now there's a weird NaN. The wonders never cease

app/Main.hs

@@ -1,15 +1,20 @@
 module Main where
 
 import Control.Lens
-import Criterion.Main
+import Control.Lens.Unsound (adjoin)
+import Control.Monad.Writer
+-- import Criterion.Main
+import Data.Default
 import Data.Either
 import Petzval.Optics
 import Petzval.Optics.RTM
 import Petzval.System
+import Petzval.Merit
 import Petzval.Trace
 import Petzval.Calculations
 import Petzval.Optimization
 import Petzval.Types
+import System.Environment (getArgs)
 import Linear
 import Numeric.AD.Mode (Scalar, Mode)
 import qualified Numeric.LinearAlgebra as L
@@ -41,13 +46,13 @@ system2 =
   , 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
+  where s_sk16 = constMat 1.620408 :: SellemeierMat
-        s_f4 = ConstMat 1.616589
+        s_f4 = constMat 1.616589
 
 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
+  where baked = bake 0.5875618 system
         ep = entrancePupil baked
         result = sum
           . map (^2)
@@ -63,6 +68,17 @@ merit system = result
 type Id a = a -> a
 
 
+doOptimize steps = runWriterT $ optimizeDLS cfg vars merit system1
+    where merit = mkMerit [ DynMerit $ TW 50 100 $ BFL Nothing
+                          , DynMerit $ TW 0 1 $ SpotSize Nothing
+                          ]
+          vars :: VariableSet
+          vars = (ix 1 . roc) `adjoin` (ix 2 . roc)
+          vars2 = undefined -- foldl1 adjoin
+                  -- [ ix 1.roc , ix 2.roc]
+          cfg = def{maxSteps = steps}
+
+              
 
 main :: IO ()
 main =
@@ -70,10 +86,14 @@ main =
       ep = (entrancePupil baked){position=20.4747094540}
       fa = 20
   in do
+    [arg] <- getArgs
+    let steps = read arg
     print $ L.norm_2  (L.fromList [3.0 , 4.0] :: L.Vector Double)
+    doOptimize steps >>= print
     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])

lib/Petzval/Calculations.hs

@@ -15,7 +15,8 @@ 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'
+rmsSize [] = 1/0
+rmsSize points =  rmsSize' (centroid points) points
 
 centroid :: Fractional a => [V2 a] -> V2 a
 centroid = uncurry (^/) . foldl1' (bimap2 (^+^) (+)) . map (flip (,) 1)

lib/Petzval/Merit.hs

@@ -1,33 +1,123 @@
-module Petzval.Merit where
+{- LANGUAGE ImpredicativeTypes #-}
+module Petzval.Merit (BFL(..), SpotSize(..), DynMerit(..), TW(..), Merit(..), MeritFunction, evalMerit, mkMerit) where
 
 import Petzval.Calculations
 import Petzval.Types
 import Petzval.Trace
-import Petzval.Optics (Element, BakedIOR)
+import Petzval.Optics (Element, BakedIOR, SellemeierMat, bake)
 import Petzval.System
 import Control.Lens
 import Data.Either
+import Data.Default
+import Petzval.Optics.RTM
 import Linear
+import Data.Maybe
+import Numeric.AD.Mode (auto)
 
 import qualified Data.Map as Map
 
-type MeritPart a = [Element BakedIOR a] -> a
+type MeritPart a = TracedSystem a -> a
 
-data Calcuable a => TracedSystem a =
+
-  TracedSystem { system :: [Element BakedIOR a]
+data TraceConditions = TraceConditions { fieldAngle :: Double
-               , field_angles :: [Double]
+                                       , wavelength :: Double
-               , wavelengths :: [Double]
-               , tracePoints :: [(Ray a, HitRecord a)]
                                        }
+                     deriving (Show)
+    
+data TracedSystem a =
+  TracedSystem { _innerSystem :: [Element SellemeierMat a]
+               , _fieldAngles :: [Double]
+               , _wavelengths :: [Double]
+               --, _tracePoints :: Map.Map TraceConditions [(Ray a, HitRecord a)]
+               }
+makeLenses ''TracedSystem
+
+defWavelength :: (TracedSystem a) -> Double
+defWavelength = fromMaybe 0.5875618 . listToMaybe  . _wavelengths
+defFA :: TracedSystem a -> Double
+defFA = fromMaybe 0 . listToMaybe . _fieldAngles
+defTraceConditions :: TracedSystem a -> TraceConditions
+defTraceConditions = TraceConditions <$> defFA <*> defWavelength
+
+systemAtWavelength :: Double -> TracedSystem a -> [Element BakedIOR a]
+systemAtWavelength λ =  bake λ . _innerSystem
+
+class Merit m where
+    calc :: Calcuable a => m -> TracedSystem a -> a
 
 
+data SpotSize = SpotSize (Maybe TraceConditions)
+                deriving (Show)
 
-spotSize :: Calcuable a => a -> MeritPart a
+instance Merit SpotSize where
-spotSize fa system = rmsSize
+    calc (SpotSize Nothing) sys = calc (SpotSize (Just $ defTraceConditions sys)) sys
+    calc (SpotSize (Just a)) sys = spotSize' a sys
+
+data BFL = BFL (Maybe Double)
+         deriving (Show)
+instance Merit BFL where
+    calc (BFL Nothing) = bfl
+    calc (BFL (Just wl)) = bfl' wl
+                         
+spotSize' :: Calcuable a => TraceConditions -> MeritPart a
+spotSize' conditions ts = rmsSize
                              . toListOf (each._pos._xy)
                              . rights
                              . map fst
                              . map (raytrace system)
                              . map (createRay Nothing ep fa)
                              $ hexapolarPattern 10
-  where ep = entrancePupil system
+    where
+      system = systemAtWavelength (wavelength conditions) ts
+      fa = auto $ fieldAngle conditions
+      ep = entrancePupil system
+
+spotSize :: Calcuable a => MeritPart a
+spotSize = defTraceConditions >>= spotSize'
+
+-- | Compute the back focal length
+bfl :: Calcuable a => MeritPart a
+-- | Compute the back focal length at a specific wavelength
+bfl' :: (Calcuable a)
+        => Double -- ^ Wavelength
+     -> MeritPart a -- ^ Merit function
+
+bfl' wavelength = rearFocalLength . systemRTM . systemAtWavelength wavelength
+bfl = defWavelength >>= bfl'
+
+
+data TW m = TW Double Double m
+          deriving (Show)
+instance Merit m => Merit (TW m) where
+    calc (TW target weight merit) system = auto weight * (calc merit system - auto target)
+
+component :: Calcuable a => Double -> Double -> MeritPart a -> MeritPart a
+component target weight part system = let value = part system
+                                      in (value - auto target) * auto weight
+
+
+data MeritFunction = MeritFunction { components :: [DynMerit]
+                                   , fnFieldAngles :: [Double]
+                                   , fnWavelengths :: [Double]
+                                   }
+mkMerit :: [DynMerit] -> MeritFunction
+mkMerit parts =  MeritFunction {
+                   components = parts
+                 , fnFieldAngles = []
+                 , fnWavelengths = []
+                 }
+
+data DynMerit = forall m. (Show m, Merit m) => DynMerit m
+instance Merit DynMerit where
+    calc (DynMerit m) = calc m
+                   
+evalMerit :: Calcuable a => MeritFunction -> [Element SellemeierMat a] -> [a]
+evalMerit fn system =
+    let ts = TracedSystem {
+               _innerSystem=system
+             , _fieldAngles=fnFieldAngles fn
+             , _wavelengths = fnWavelengths fn
+             }
+    in fmap (flip calc ts) $ components fn
+
+      -- fmap ($ts) $ components fn

lib/Petzval/Optics.hs

@@ -90,9 +90,9 @@ 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 =
+bake wavelength system =
   snd $
   mapAccumLOf (each.material) (\n1 (ConstMat n2) -> (n2, BakedIOR n1 n2)) 1 $
-  specialize wavelength mat
+  specialize wavelength system
 
 

lib/Petzval/Optimization.hs

@@ -2,6 +2,7 @@
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE NoMonomorphismRestriction #-}
+{- LANGUAGE ImpredicativeTypes #-}
 module Petzval.Optimization where
 
 import Linear
@@ -13,11 +14,15 @@ import Control.Lens
 import Control.Lens.Unsound (adjoin)
 import Numeric.AD.Mode (auto)
 import Numeric.AD.Mode.Reverse.Double
+import Petzval.Merit
 import Petzval.Optics
 import Petzval.Types
 import Control.Monad.Writer.Class
+import Data.Default
 
-import Numeric.LinearAlgebra hiding (Element, dot)
+import Debug.Trace
+
+import Numeric.LinearAlgebra hiding (Element, dot, (??))
 import qualified Numeric.LinearAlgebra as L
 
 
@@ -53,23 +58,40 @@ data DLSCfg = DLSCfg { eps1 :: Double -- ^ Cutoff for the derivative of the metr
                      , eps2 :: Double -- ^ Cutoff for no longer making progress
                      , maxSteps :: Integer -- ^ Maximum number of steps to iterate
                      }
+
+instance Default DLSCfg where
+    def = DLSCfg 1e-8 1e-4 1000
+
 makeLenses ''DLSState
 
+           {-
+evalMeritJ :: MeritFunction -> VariableSet -> [Element SellemeierMat Double] -> [(Double, [Double])]
+evalMeritJ fn vars system = 
+    let ts sys = TracedSystem {
+               _system=sys
+             , _fieldAngles=fieldAngles fn
+             , _wavelengths = wavelengths fn
+             }
+        evalAt pt = fmap ($lts) $ components fn
+            where lts = ts $ setVars vars (system & each.liftFp %~ auto) pt
+    in unzip . jacobian' . evalAt $ extractVars vars system
+-}
+
 optimizeDLS :: (MonadWriter [[Double]] m -- ^ This yields a list of intermediate values of the merit function
-               )
+               , MonadIO m)
             => DLSCfg
             -> VariableSet -- ^ The set of independent variables
-            -> (forall a. Calcuable a => [Element mat a] -> [a]) -- ^ merit function
+            -> MeritFunction -- ^ merit function
-            -> [Element mat Double] -- ^ The system
+            -> [Element SellemeierMat Double] -- ^ The system
-            -> m [Element mat Double]
+            -> m [Element SellemeierMat Double]
 
 optimizeDLS cfg vars merit system = fmap (setVars vars system . toList) $ evalStateT doOptimize initialState 
   where
     initialState = let pt = fromList $ extractVars vars system
-                       (y,j) = jacobianAt pt
+                       (y,j) = traceShowId $ jacobianAt pt
                        a :: Matrix Double = tr j L.<> j
                        damping :: Double = 1e-3 * (maximum . toList . takeDiag) a
-                       lastSum = sum . map (^2) . merit $ system
+                       lastSum = sum . map (^2) . evalMerit merit $ system
                    in DLSState { _damping=damping
                                , _lastSum=lastSum
                                , _curPt=pt
@@ -82,13 +104,14 @@ optimizeDLS cfg vars merit system = fmap (setVars vars system . toList) $ evalSt
     -- optimizeStep :: m Double, where the return value is the change in merit
     optimizeStep = do
       mu <- use damping
+      liftIO . print $ "Damping: " ++ show mu
       lastPt <- use curPt
       let (y, a) = jacobianAt lastPt
       let g = tr a #> y
       let dPt :: Vector Double = -(inv $ tr a L.<> a + (scalar mu * ident (cols a)) ) #> g
       let newPt = lastPt + dPt
       let oldMerit = sumSq y
-      let newMerit = sumSq . fromList . merit . setVars vars system . toList $ newPt
+      let newMerit = sumSq . fromList . evalMerit merit . setVars vars system . toList $ newPt
       let dL = (dPt `L.dot` (scalar mu * dPt - g)) / 2
       let gain = (oldMerit - newMerit) / dL
       if gain > 0
@@ -98,14 +121,17 @@ optimizeDLS cfg vars merit system = fmap (setVars vars system . toList) $ evalSt
                 found ||= (norm_2 g <= eps1 cfg || norm_2 dPt <= eps2 cfg * (norm_2 lastPt + eps2 cfg))
                 curPt .= newPt
         else do scale <- use dampScale
+                d <- use damping
                 dampScale *= 2
+                liftIO $ print ("Rescaling", scale, d)
                 damping *= scale
         
       tell [toList y]
       
       return $ newMerit - oldMerit
     jacobianAt :: Vector Double -> (Vector Double, Matrix Double)
-    jacobianAt pt = let (y,a) = unzip . jacobian' (merit . setVars vars (system & each.liftFp %~ auto)) $ toList pt
+    jacobianAt pt = let eval pt = undefined
+                        (y,a) = unzip . jacobian' (\adpt -> evalMerit merit $ setVars vars (system & each.liftFp %~ auto) adpt) $ toList pt
                     in (fromList y, fromLists a)
 
 sumSq :: Vector Double -> Double

lib/Petzval/Types.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE FlexibleInstances, UndecidableInstances #-}
+{-# LANGUAGE FlexibleInstances, UndecidableInstances, FlexibleContexts #-}
 module Petzval.Types(Calcuable) where
 
 import Numeric.AD.Mode

petzval-hs.cabal

@@ -38,8 +38,10 @@ executable petzval-hs
                       ad,
                       linear,
                       lens,
-                      criterion,
+                      mtl,
+                      -- criterion,
                       hmatrix,
+                      data-default,
                       petzval-hs
     hs-source-dirs:   app
 
@@ -48,10 +50,12 @@ executable petzval-prof
     main-is: Main.hs
     build-depends:    base,
                       ad,
+                      mtl,
                       linear,
                       lens,
-                      criterion,
+                      -- criterion,
                       hmatrix,
+                      data-default,
                       petzval-hs
     hs-source-dirs:   app
     ghc-options:
@@ -83,4 +87,5 @@ library
                       deepseq,
                       containers,
                       monad-loops,
-                      hmatrix
+                      hmatrix,
+                      data-default