Implemented some form of LM optimization, but it's proper fucked somehow.

lib/Petzval/Optics.hs

@@ -28,6 +28,8 @@ class Material mat where
   iorAtWavelength :: Double -> mat -> Double
   -- | A material with a constant IOR of 1
   air :: mat
+  air = constMat 1
+  constMat :: Double -> mat
   
 -- | Calculates IOR based on the Sellemeier equation
 data SellemeierMat = SellemeierMat [(Double, Double)]
@@ -37,6 +39,7 @@ instance Material SellemeierMat where
     where contrib (b,c) = b * w2 / (w2 - c)
           w2 = λ ^ 2
   air = SellemeierMat []
+  constMat ior = SellemeierMat [(ior * ior - 1, 0)]
 
 -- | A material with a constant IOR
 newtype ConstMat = ConstMat Double
@@ -45,6 +48,7 @@ newtype ConstMat = ConstMat Double
 instance Material ConstMat where
   iorAtWavelength _ (ConstMat ior) = ior
   air = ConstMat 1
+  constMat = ConstMat
 
 -- | A representation of IOR as (previous,new) IOR at a surface
 data BakedIOR = BakedIOR Double Double
@@ -74,30 +78,7 @@ isStop _ = False
 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
-thickness inj (s@Surface{_thickness}) = (\nt -> s{_thickness=nt}) <$> inj _thickness
-thickness inj (s@Stop{_thickness}) = (\nt -> s{_thickness=nt}) <$> inj _thickness
--- | The outside radius of the element. Rays that intersect the element beyond this radius are considered to have missed.
-outsideRadius :: Lens (Element mat a) (Element mat a) a a
-outsideRadius inj (s@Surface{_outsideRadius}) = (\nt -> s{_outsideRadius=nt}) <$> inj _outsideRadius
-outsideRadius inj (s@Stop{_outsideRadius}) = (\nt -> s{_outsideRadius=nt}) <$> inj _outsideRadius
 
-infinity :: (RealFloat a) => a
-infinity = encodeFloat (floatRadix 0 - 1) (snd $ floatRange 0)
--- | The radius of curvature of an element. This is 1 for stops
-roc :: RealFloat a => Lens (Element mat a) (Element mat a) a a
-roc inj (s@Surface{_roc}) = (\nt -> s{_roc=nt}) <$> inj _roc
-roc inj (s@Stop{}) = (const s) <$> inj infinity
-
--- | The material that a surface transitions into
---
--- For stops, this returns None and ignores being set.
-material :: (Material mat, Material mat') => Lens (Element mat a) (Element mat' a) (Maybe mat) (Maybe mat')
-material inj (s@Surface{_thickness, _outsideRadius, _roc, _material}) = (\nt -> Surface{_material=fromMaybe air nt, _thickness, _outsideRadius, _roc}) <$> inj (Just _material)
-material inj (s@Stop{_thickness, _outsideRadius}) = (const Stop{_thickness, _outsideRadius}) <$> inj Nothing
--}
 -- | 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@Stop{_thickness, _outsideRadius}) = Stop <$> inj _thickness <*> inj _outsideRadius

lib/Petzval/Optimization.hs

@@ -1,12 +1,25 @@
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE NoMonomorphismRestriction #-}
 module Petzval.Optimization where
 
+import Linear
+import Control.Monad
+import Control.Monad.Loops
+import Control.Monad.State
+import Control.Monad.Writer (MonadWriter, tell)
 import Control.Lens
 import Control.Lens.Unsound (adjoin)
-import Numeric.AD.Mode
+import Numeric.AD.Mode (auto)
 import Numeric.AD.Mode.Reverse.Double
 import Petzval.Optics
 import Petzval.Types
+import Control.Monad.Writer.Class
+
+import Numeric.LinearAlgebra hiding (Element, dot)
+import qualified Numeric.LinearAlgebra as L
+
 
 -- | 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:
@@ -28,6 +41,73 @@ gradientAt :: VariableSet -- ^ The set of independent variables
            -> (Double, [Double]) -- ^ The gradient
 gradientAt vars merit system =  grad' (merit . setVars vars (system & each.liftFp %~ auto)) (extractVars vars system)
            
+data DLSState = DLSState { _damping :: Double
+                         , _lastSum :: Double
+                         , _curPt :: Vector Double
+                         , _dampScale :: Double
+                         , _found :: Bool
+                         , _curIter :: Integer
+                         }
 
+data DLSCfg = DLSCfg { eps1 :: Double -- ^ Cutoff for the derivative of the metric function
+                     , eps2 :: Double -- ^ Cutoff for no longer making progress
+                     , maxSteps :: Integer -- ^ Maximum number of steps to iterate
+                     }
+makeLenses ''DLSState
 
+optimizeDLS :: (MonadWriter [[Double]] m -- ^ This yields a list of intermediate values of the merit function
+               )
+            => DLSCfg
+            -> VariableSet -- ^ The set of independent variables
+            -> (forall a. Calcuable a => [Element mat a] -> [a]) -- ^ merit function
+            -> [Element mat Double] -- ^ The system
+            -> m [Element mat 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
+                       a :: Matrix Double = tr j L.<> j
+                       damping :: Double = 1e-3 * (maximum . toList . takeDiag) a
+                       lastSum = sum . map (^2) . merit $ system
+                   in DLSState { _damping=damping
+                               , _lastSum=lastSum
+                               , _curPt=pt
+                               , _dampScale = 2
+                               , _found = False
+                               , _curIter = maxSteps cfg
+                               }
+    isDone = orM [use found, uses curIter (<0)]
+    doOptimize = (untilM_ (curIter -= 1 >> optimizeStep) isDone) >> use curPt
+    -- optimizeStep :: m Double, where the return value is the change in merit
+    optimizeStep = do
+      mu <- use damping
+      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 dL = (dPt `L.dot` (scalar mu * dPt - g)) / 2
+      let gain = (oldMerit - newMerit) / dL
+      if gain > 0
+        then do curPt .= lastPt `add` dPt
+                damping .= mu * max (1/3) (1 - (2*gain - 1) ^ 3)
+                dampScale .= 2
+                found ||= (norm_2 g <= eps1 cfg || norm_2 dPt <= eps2 cfg * (norm_2 lastPt + eps2 cfg))
+                curPt .= newPt
+        else do scale <- use dampScale
+                dampScale *= 2
+                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
+                    in (fromList y, fromLists a)
+
+sumSq :: Vector Double -> Double
+sumSq x = L.dot x x
 -- instances

lib/Petzval/System.hs

@@ -59,17 +59,7 @@ instance Num a => Semigroup (Seidel a) where
 instance Num a => Monoid (Seidel a) where
   mempty = Seidel { sphr=0, coma=0, asti=0, fcur=0, dist=0
                   -- , c1 = mempty, c2 = mempty
-                  }
-  mappend Seidel{sphr, coma, asti, fcur, dist} Seidel{sphr=sphr', coma=coma', asti=asti', fcur=fcur', dist=dist'} =
-    Seidel { sphr = sphr + sphr'
-           , coma = coma + coma'
-           , asti = asti + asti'
-           , fcur = fcur + fcur'
-           , dist = dist + dist'
-           --, c1 = c1 + c1'
-           --, c2 = c2 + c2'
-           }
-    
+                  }    
 
 
 -- | Initial matrix is [ h_ h; u_ u ]

lib/Petzval/Trace.hs

@@ -142,8 +142,8 @@ raytrace system ray =
 raytrace1 :: ( Floating a, Ord a, Mode a, Scalar a ~ Double, Epsilon a
              , MonadState Double m
              , MonadWriter [HitRecord a] m -- ^ Tracing yields a list of 
-             , MonadError TraceError m) => -- ^ This can fail
-             Ray a -> Element BakedIOR a -> m (Ray a)
+             , MonadError TraceError m)  -- ^ This can fail
+          => Ray a -> Element BakedIOR a -> m (Ray a)
 raytrace1 ray element = do
       n1 <- get
       let stopP = isStop element