RANSAC-0.1.0.0: tests/LinearFit.hs

-- | Example use of the RANSAC algorithm to fit a line to some
-- points. We start with points generated by a process defined by the
-- equation of a line in 2D. These points are affected by normally
-- distributed noise, and our data set is further corrupted by a "red
-- herring" cluster of points that we would like to ignore. We use
-- RANSAC to cut through the noise and fit a line to the point data
-- set.
-- 
-- The important feature of RANSAC as applied here is that it manages
-- to ignore the spurious (red herring) cluster centered at (0,8).
--
-- The Chart package is used to visualize the data and estimated
-- model.
module Main where
import Control.Applicative
import Control.Lens (view)
import Data.Accessor ((^=))
import Data.Colour (opaque)
import Data.Colour.Names
import qualified Data.Foldable as F
import Data.Random.Normal (normalsIO')
import Data.Vector.Storable (Vector)
import qualified Data.Vector.Storable as V
import Graphics.Rendering.Chart hiding (Vector,Point)
import Linear
import Numeric.Ransac

type Point = V2 Float

-- | Fit a 2D line to a collection of 'Point's.
fitLine :: Vector Point -> Maybe (V2 Float)
fitLine pts = (!* b) <$> inv22 a
  where sx = V.sum $ V.map (view _x) pts
        a = V2 (V2 (V.sum (V.map ((^2).view _x) pts)) sx)
               (V2 sx (fromIntegral (V.length pts)))
        b = V2 (V.sum (V.map F.product pts))
               (V.sum (V.map (view _y) pts))

-- | Compute the error of a 'Point' with respect to a hypothesized
-- linear model.
ptError :: V2 Float -> Point -> Float
ptError (V2 m b) (V2 x y) = sq $ y - (m*x+b)
  where sq x = x * x

-- | Produce a plot of all the points we have to work with. A green
-- dashed line indicates the ground truth linear model, the solid
-- purple line shows the RANSAC model, and the points that are inliers
-- for that model are circled in yellow.
main = do noise <- v2Cast . V.fromList . take (n*2) <$> normalsIO' (0,0.3)
          herring <- V.zipWith V2 
                     <$> (V.fromList . take 200 <$> normalsIO' (0,0.2))
                     <*> (V.fromList . take 200 <$> normalsIO' (8,0.6))
          let pts' = V.zipWith (+) noise pts
          let pts'' = pts' V.++ herring
          res <- ransac 100 2 0.5 fitLine ptError (< 2) pts''
          case res of
            Nothing -> putStrLn "No model found"
            Just (model,inliers) -> 
              do putStrLn $ "Model "++show model++" with "++
                            show (V.length inliers)++" inliers"
                 let pp = PlotPoints "data" 
                                     (filledCircles 2 (opaque blue))
                                     (map (toTup . dub) (V.toList pts''))
                     ppi = PlotPoints "inliers"
                                      (hollowCircles 3 2 (opaque yellow))
                                      (map (toTup . dub) (V.toList inliers))
                     lp = PlotLines "truth"
                                    (dashedLine 3 [10,10] (opaque green))
                                    [[ toTup $ dub (mkPt 0) 
                                     , toTup $ dub (mkPt (n-1)) ]]
                                    []
                     lp' = PlotLines "model"
                                     (solidLine 5 (opaque purple))
                                     [[ toTup $ dub (mkPt' model 0)
                                      , toTup $ dub (mkPt' model (n-1)) ]]
                                     []
                     layout = layout1_title ^="2D Linear Fit"
                            $ layout1_background ^= solidFillStyle (opaque white)
                            $ layout1_plots ^= [ Left (toPlot pp)
                                               , Left (toPlot ppi)
                                               , Left (toPlot lp)
                                               , Left (toPlot lp') ]
                            $ setLayout1Foreground (opaque black)
                            $ defaultLayout1
                 renderableToPDFFile (toRenderable layout) 600 600 "foo.pdf"
  where n = 1000
        pts = V.generate n mkPt
        mkPt :: Int -> V2 Float
        mkPt i = let x = fromIntegral i / 500
                 in V2 x (5*x + 2)
        v2Cast :: Vector Float -> Vector Point
        v2Cast = V.unsafeCast
        toTup (V2 x y) = (x,y)
        dub :: V2 Float -> V2 Double
        dub = fmap realToFrac
        mkPt' (V2 m b) i = let x = fromIntegral i / 500
                           in V2 x (x * m + b)