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

2023-01-30 02:07:10 +01:00
parent cb2dad39ec
commit 32fb672bf5
7 changed files with 107 additions and 42 deletions
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -12,6 +12,7 @@ import Petzval.Optimization
 import Petzval.Types
 import Linear
 import Numeric.AD.Mode (Scalar, Mode)
+import qualified Numeric.LinearAlgebra as L

 bk7 = SellemeierMat [ (1.03961212,  6.00069867e-3 )
                    , (0.231792344, 2.00179144e-2 )
@@ -56,7 +57,7 @@ merit system = result
                  . rights
                  . map fst
                  . map (raytrace baked . createRay Nothing ep angle)
-                  $ spiralPattern 10)
+                  $ hexapolarPattern 6)
          $ [0, 7.07, 15]

 type Id a = a -> a
@@ -69,13 +70,17 @@ main =
      ep = (entrancePupil baked){position=20.4747094540}
      fa = 20
  in do
+    print $ L.norm_2  (L.fromList [3.0 , 4.0] :: L.Vector Double)
    putStrLn $ show ep
    putStrLn $ show $ rearFocalLength $ lensRTM baked
    putStrLn $ show $ mconcat (seidel ep fa baked :: [Seidel Double])
-    defaultMain [
+ {-   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
+      , bgroup "ad" [ bench "system1" $ nf (gradientAt (ix 1.thickness) merit) system1,
+                      bench "system2" $ nf (gradientAt (ix 1.thickness) merit) system2
+
                    ]
      ]
+-}
--- a/flake.nix
+++ b/flake.nix
@@ -15,13 +15,17 @@
          hello = flake-utils.lib.mkApp { drv = self.packages.${system}.hello; };
          default = hello;
        };
-        devShell = pkgs.mkShell {
+        devShell = pkgs.mkShell rec {
          buildInputs = with pkgs; [
            cabal-install
            stack
            haskell-language-server
            ghc
+
+            openblas
+            openblas.out
          ];
+          LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath buildInputs;
        };
      }
    );
--- a/lib/Petzval/Optics.hs
+++ b/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
--- a/lib/Petzval/Optimization.hs
+++ b/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
--- a/lib/Petzval/System.hs
+++ b/lib/Petzval/System.hs
@@ -60,16 +60,6 @@ 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 ]
--- a/lib/Petzval/Trace.hs
+++ b/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
--- a/petzval-hs.cabal
+++ b/petzval-hs.cabal
@@ -39,6 +39,7 @@ executable petzval-hs
                      linear ^>=1.21,
                      lens   ^>=5.0,
                      criterion ^>=1.5,
+                      hmatrix ^>= 0.20.2,
                      petzval-hs
    hs-source-dirs:   app

@@ -68,6 +69,7 @@ library
        Petzval.Calculations
        Petzval.Optimization
        Petzval.Types
+        Petzval.Merit
    other-modules:
        Petzval.Internal.Vec
    hs-source-dirs: lib
@@ -77,4 +79,7 @@ library
                      linear ^>=1.21,
                      reflection ^>=2.1,
                      mtl    ^>=2.2,
-                      deepseq ^>=1.4
+                      deepseq ^>=1.4,
+                      containers ^>=0.6.4.1,
+                      monad-loops ^>= 0.4.3,
+                      hmatrix ^>= 0.20.2