packages feed

brush-strokes-0.1.0.0: src/lib/Math/Bezier/Cubic/Fit.hs

{-# LANGUAGE ScopedTypeVariables #-}

module Math.Bezier.Cubic.Fit
  ( FitParameters(..), FitPoint(..), FwdBwd(..)
  , OutlinePoint(..)
  , fitSpline, fitPiece
  )
  where

-- base
import Control.Arrow
  ( first, second )
import Control.Monad
  ( when )
import Control.Monad.ST
  ( ST, runST )
import Data.Complex
  ( Complex(..) )
import Data.Foldable
  ( for_ )
import Data.Functor
  ( ($>) )
import Data.Semigroup
  ( Arg(..), Max(..), ArgMax )
import GHC.Generics
  ( Generic )

-- acts
import Data.Act
  ( Act
    ( (•) )
  )

-- containers
import Data.Sequence
  ( Seq(..) )
import qualified Data.Sequence as Seq
  ( fromList )

-- deepseq
import Control.DeepSeq
  ( NFData )

-- parallel
import qualified Control.Parallel.Strategies as Parallel.Strategy
  ( rdeepseq, parTuple2, using )

-- primitive
import Data.Primitive.PrimArray
  ( MutablePrimArray
  , generatePrimArray
  , primArrayFromListN
  , unsafeThawPrimArray, unsafeFreezePrimArray
  , readPrimArray, writePrimArray, primArrayToList
  )

-- transformers
import Control.Monad.Trans.State.Strict
  ( execStateT, modify' )
import Control.Monad.Trans.Class
  ( lift )

-- brush-strokes
import qualified Math.Bezier.Cubic as Cubic
  ( Bezier(..), bezier, ddist )
import Math.Bezier.Spline
  ( SplineType(..), SplinePts
  , openCubicBezierCurveSpline
  )
import Math.Epsilon
  ( epsilon )
import Math.Linear.Solve
  ( linearSolve )
import Math.Module
  ( Module((*^), (^-^)), lerp
  , Inner((^.^)), quadrance
  )
import Math.Roots
  ( laguerre ) --, eval, derivative )
import Math.Linear
  ( Mat22(..), ℝ(..), T(..) )

--------------------------------------------------------------------------------

-- | Parameters to the curve fitting algorithm.
data FitParameters
  = FitParameters
    { maxSubdiv  :: !Int    -- ^ maximum subdivision recursion number
    , nbSegments :: !Int    -- ^ number of segments to split the curve into for fitting
    , dist_tol   :: !Double -- ^ tolerance for the distance
    , t_tol      :: !Double -- ^ the tolerance for the Bézier parameter (for the fitting process)
    , maxIters   :: !Int    -- ^ maximum number of iterations (for the fitting process)
    }
    deriving stock Show

data FwdBwd
  = Fwd
  | Bwd
  deriving stock ( Eq, Show, Generic )
  deriving anyclass NFData

-- | Point computed as being on the outline.
--
-- Output of the envelope equation calculation, and input to the
-- curve fitting algorithm.
data OutlinePoint
  = OutlinePoint
  { outlinePoint      :: ℝ 2
  , outlineParam      :: Double
  , outlineTangent    :: T ( ℝ 2 )
  , outlineDotProduct :: Double
  }
  deriving stock ( Eq, Show, Generic )

data FitPoint
  = FitPoint
    { fitDir        :: !FwdBwd
    , fitPoint      :: !( ℝ 2 )
    , fitParam      :: !( ℝ 1, ℝ 1 )
    , fitDotProduct :: !Double
    }
  | FitTangent
    { fitDir        :: !FwdBwd
    , fitPoint      :: !( ℝ 2 )
    , fitTangent    :: !( T ( ℝ 2 ) )
    , fitDotProduct :: !Double
    }
  | JoinPoint
    { joinDir       :: !FwdBwd
    , joinPoint     :: !( ℝ 2 )
    }
  deriving stock    ( Show, Generic )
  deriving anyclass NFData

-- | Fits a cubic Bézier spline to the given curve \( t \mapsto C(t), 0 \leqslant t \leqslant 1 \),
-- assumed to be G1-continuous.
--
-- Additionally returns the points that were used to perform the fit, for debugging purposes.
--
-- Subdivides the given curve into the specified number of segments \( \texttt{nbSegments} \),
-- and tries to fit the resulting points with a cubic Bézier curve using 'fitPiece'.
--
-- When the fit is too poor: subdivide at the worst fitting point, and recurse.
-- \( \texttt{maxSubdiv} \) controls the maximum recursion depth for subdivision.
--
-- See 'fitPiece' for more information on the fitting process,
-- including the meaning of \( \texttt{t_tol} \) and \( \texttt{maxIters} \).
fitSpline
  :: FwdBwd
  -> FitParameters
  -> ( ℝ 1 -> OutlinePoint ) -- ^ curve \( t \mapsto ( C(t), C'(t) ) \) to fit
  -> ( ℝ 1, ℝ 1 )            -- ^ interval \( [t_min, t_max] \)
  -> ( SplinePts Open, Seq FitPoint )
fitSpline fwdOrBwd ( FitParameters {..} ) curveFn = go 0
  where
    dt :: Double
    dt = recip ( fromIntegral nbSegments )
    go
      :: Int
      -> ( ℝ 1, ℝ 1 )
      -> ( SplinePts Open, Seq FitPoint )
    go subdiv ( t_min, t_max ) =
      let
        p, r :: ℝ 2
        tp, tr :: T ( ℝ 2 )
        qs :: [ ( ℝ 1, ( ℝ 2, Double ) ) ]
        OutlinePoint { outlinePoint = p, outlineTangent = tp, outlineDotProduct = dot_p } = curveFn t_min
        OutlinePoint { outlinePoint = r, outlineTangent = tr, outlineDotProduct = dot_r } = curveFn t_max
        qs = [ ( t0, ( q, dot_q ) )
             | j <- [ 1 .. nbSegments - 1 ]
             , let t0 = lerp @( T ( ℝ 1 ) ) ( dt * fromIntegral j ) t_min t_max
                   OutlinePoint { outlinePoint = q, outlineDotProduct = dot_q }
                     = curveFn t0
             ]
      in
        case fitPiece dist_tol t_tol maxIters p tp ( fmap ( fst . snd ) qs ) r tr of
          ( bez, ts', Max ( Arg max_sq_error t_split_0 ) )
            |  subdiv >= maxSubdiv
            || max_sq_error <= dist_tol ^ ( 2 :: Int )
            -> -- trace ( unlines [ "fitSpline: piece is OK", "t_min = " ++ show t_min, "start = " ++ show p, "start tgt = " ++ show tp, "t_max = " ++ show t_max, "end = " ++ show r, "end tgt = " ++ show tr ] ) $
               ( openCubicBezierCurveSpline () bez
               ,    ( FitTangent fwdOrBwd p tp dot_p
                 :<| Seq.fromList ( zipWith ( \ ( t0, ( pt, dot ) ) δ -> FitPoint fwdOrBwd pt ( t0, lerp @( T ( ℝ 1 ) ) δ t_min t_max ) dot ) qs ts' ) )
                 :|> FitTangent fwdOrBwd r tr dot_r
               )
            | let
                t_split :: ℝ 1
                t_split = ℝ1 $ min ( unℝ1 $ lerp @( T ( ℝ 1 ) ) ( dt * 1 )                               t_min t_max )
                             $ max ( unℝ1 $ lerp @( T ( ℝ 1 ) ) ( dt * fromIntegral ( nbSegments - 1 ) ) t_min t_max )
                             $ t_split_0
                c1, c2 :: SplinePts Open
                ps1, ps2 :: Seq FitPoint
                ( ( c1, ps1 ), ( c2, ps2 ) )
                  = ( go ( subdiv + 1 ) ( t_min  , t_split )
                    , go ( subdiv + 1 ) ( t_split, t_max   )
                    ) `Parallel.Strategy.using`
                       ( Parallel.Strategy.parTuple2 Parallel.Strategy.rdeepseq Parallel.Strategy.rdeepseq )
            -> ( c1 <> c2, ps1 <> ps2 )

-- | Fits a single cubic Bézier curve to the given data.
--
-- This will consist of a cubic Bézier curve which:
--
--   * starts at \( p \) with tangent \( \textrm{t}_p \),
--   * ends at \( r \) with tangent \( \textrm{t}_r \),
--   * best fits the intermediate sequence of points \( \left ( q_i \right )_{i=1}^n \).
--
-- This function also returns \( \textrm{ArgMax}\ d^2_\textrm{max}\ t_\textrm{max}: \)
-- the parameter and squared distance of the worst-fitting point.
-- It is guaranteed that all points to fit lie within the tubular neighbourhood
-- of radius \( d_\textrm{max} \) of the fitted curve.
--
-- /Note/: the order of the intermediate points is important.
--
-- Proceeds by fitting a cubic Bézier curve \( B(t) \), \( 0 \leqslant t \leqslant 1 \),
-- with given endpoints and tangents, which minimises the sum of squares functional
--
-- \[ \sum_{i=1}^n \Big \| B(t_i) - q_i \Big \|^2. \]
--
-- The values of the parameters \( \left ( t_i \right )_{i=1}^n \) are recursively estimated,
-- starting from uniform parametrisation (this will be the fit when \( \texttt{maxIters} = 0 \)).
--
-- The iteration ends when any of the following conditions are satisfied:
--
--   * each new estimated parameter value \( t_i' \) differs from
--     its previous value \( t_i \) by less than \( \texttt{t_tol} \),
--   * each on-curve point \( B(t_i) \) is within distance \( \texttt{dist_tol} \)
--     of its corresponding point to fit \( q_i \),
--   * the maximum iteration limit \( \texttt{maxIters} \) has been reached.
fitPiece
  :: Double    -- ^ \( \texttt{dist_tol} \), tolerance for the distance
  -> Double    -- ^ \( \texttt{t_tol} \), the tolerance for the Bézier parameter
  -> Int       -- ^ \( \texttt{maxIters} \), maximum number of iterations
  -> ℝ 2       -- ^ \( p \), start point
  -> T ( ℝ 2 ) -- ^ \( \textrm{t}_p \), start tangent vector (length is ignored)
  -> [ ℝ 2 ]   -- ^ \( \left ( q_i \right )_{i=1}^n \), points to fit
  -> ℝ 2       -- ^ \( r \), end point
  -> T ( ℝ 2 ) -- ^ \( \textrm{t}_r \), end tangent vector (length is ignored)
  -> ( Cubic.Bezier ( ℝ 2 ), [ Double ], ArgMax Double Double )
fitPiece dist_tol t_tol maxIters p tp qs r tr =
  runST do
    -- Initialise the parameter values to a uniform subdivision.
    ts <- unsafeThawPrimArray $ generatePrimArray n uniform
    loop ts 0
  where
    n :: Int
    n = length qs
    uniform :: Int -> Double
    uniform i = fromIntegral ( i + 1 ) / fromIntegral ( n + 1 )

    f0, f1, f2, f3 :: Double -> T ( ℝ 2 )
    f0 t = h1 t *^ tp
    f1 t = h2 t *^ tr
    f2 t = h0 t *^ ( T p )
    f3 t = h3 t *^ ( T r )

    loop :: forall s. MutablePrimArray s Double -> Int -> ST s ( Cubic.Bezier ( ℝ 2 ), [ Double ], ArgMax Double Double )
    loop ts count = do
      let
        hermiteParameters :: Mat22 -> T ( ℝ 2 ) -> Int -> [ ℝ 2 ] -> ST s ( T ( ℝ 2 ) )
        hermiteParameters ( Mat22 a11 a12 _ a22 ) ( V2 b1 b2 ) i ( q : rest ) = do
          ti <- readPrimArray ts i
          let
            f0i, f1i, f2i, f3i :: T ( ℝ 2 )
            f0i = f0 ti
            f1i = f1 ti
            f2i = f2 ti
            f3i = f3 ti
            q'  = T q ^-^ f2i ^-^ f3i
            a11', a12', a21', a22', b1', b2' :: Double
            a11' = a11 + ( f0i ^.^ f0i )
            a12' = a12 + ( f1i ^.^ f0i )
            a21' = a12'
            a22' = a22 + ( f1i ^.^ f1i )
            b1'  = b1  + ( q'  ^.^ f0i )
            b2'  = b2  + ( q'  ^.^ f1i )
          hermiteParameters ( Mat22 a11' a12' a21' a22' ) ( V2 b1' b2' ) ( i + 1 ) rest
        hermiteParameters a b _ [] = pure ( linearSolve a b )

      ~(V2 s1 s2) <- hermiteParameters ( Mat22 0 0 0 0 ) ( V2 0 0 ) 0 qs

      let
        -- Convert from Hermite form to Bézier form.
        cp1, cp2 :: ℝ 2
        cp1 = ( (  s1 / 3 ) *^ tp ) • p
        cp2 = ( ( -s2 / 3 ) *^ tr ) • r

        bez :: Cubic.Bezier ( ℝ 2 )
        bez = Cubic.Bezier p cp1 cp2 r

      -- Run one iteration of Laguerre's method to improve the parameter values t_i,
      -- so that t_i' is a better approximation of the parameter
      -- at which the curve is closest to q_i.
      ( dts_changed, argmax_sq_dist ) <- ( `execStateT` ( False, Max ( Arg 0 0 ) ) ) $ for_ ( zip qs [ 0 .. ] ) \( q, i ) -> do
        ti <- lift ( readPrimArray ts i )
        let
          laguerreStepResult :: Complex Double
          laguerreStepResult = runST do
            coeffs <- unsafeThawPrimArray . primArrayFromListN 6
                    $ Cubic.ddist @( T ( ℝ 2 ) ) bez q
            laguerre epsilon 1 coeffs ( ti :+ 0 )
        ti' <- case laguerreStepResult of
          x :+ y
            |  isNaN x
            || isNaN y
            || abs y > epsilon
            || x < -epsilon
            || x > 1 + epsilon
            -> modify' ( first ( const True ) ) $> ti
            | otherwise
            ->  when ( abs ( x - ti ) > t_tol )
                  ( modify' ( first ( const True ) ) )
             $> ( min 1 $ max 0 x )
        let
          sq_dist :: Double
          sq_dist = quadrance @( T ( ℝ 2 ) ) q ( Cubic.bezier @( T ( ℝ 2 ) ) bez ti' )
        modify' ( second ( <> Max ( Arg sq_dist ti' ) ) )
        lift ( writePrimArray ts i ti' )

      case argmax_sq_dist of
        Max ( Arg max_sq_dist _ )
          |  count < maxIters
          && ( dts_changed || max_sq_dist > dist_tol ^ ( 2 :: Int ) )
          -> loop ts ( count + 1 )
        _ -> do
          tsList <- primArrayToList <$> unsafeFreezePrimArray ts
          pure ( bez, tsList, argmax_sq_dist )

-- | Cubic Hermite polynomial.
h0, h1, h2, h3 :: Num t => t -> t
h0 t = ( 1 + 2 * t )    * ( t - 1 ) ^ ( 2 :: Int )
h1 t = t                * ( t - 1 ) ^ ( 2 :: Int )
h2 t = t ^ ( 2 :: Int ) * ( t - 1 )
h3 t = t ^ ( 2 :: Int ) * ( 3 - 2 * t )