TestArcLength.hs

{-# OPTIONS_GHC -Wno-x-partial #-}

module Main where

-- base
import Control.Applicative
  ( (<|>) )
import Control.Monad
  ( guard )
import Data.Char
  ( toLower )
import Data.List
  ( intercalate, maximumBy )
import qualified Data.List.NonEmpty as NE
import Data.Maybe
  ( mapMaybe )
import Data.Ord
  ( comparing )
import GHC.Generics
  ( Generic )
import System.Environment
  ( withArgs )

-- deepseq
import Control.DeepSeq
  ( NFData )

-- parallel
import Control.Parallel.Strategies
  ( parListChunk, rdeepseq, withStrategy )

-- code-page
import System.IO.CodePage
  ( withCP65001 )

-- optparse-applicative
import qualified Options.Applicative as OptParse

-- falsify
import Test.Tasty.Falsify
import Test.Falsify.Predicate
  ( (.$) )
import qualified Test.Falsify.Generator as Gen
import qualified Test.Falsify.Predicate as Falsify.Prop
import qualified Test.Falsify.Range     as Range
import qualified Test.Tasty.Falsify     as Falsify

-- tasty
import qualified Test.Tasty as Tasty

-- brush-strokes
import Math.Bezier.ArcLength
  ( ArcLengthOptions(..), defaultArcLengthOptions
  , Integrator(..)
  , defaultRegulariser
  , ArcLengthParametrisation(..), curveArcLengthParametrisation
  )
import qualified Math.Bezier.Cubic     as Cubic     ( Bezier(..), bezier )
import qualified Math.Bezier.Quadratic as Quadratic ( Bezier(..), bezier )
import Math.Bezier.Spline
  ( Curve(..), NextPoint(..) )
import Math.Linear
  ( ℝ(..), T )
import Math.Module
  ( distance )

--------------------------------------------------------------------------------
-- Generator
--------------------------------------------------------------------------------

-- | Generate a 'Double' uniformly in @[lo, hi]@ on a grid of 10⁶ steps.
genDouble :: Double -> Double -> Gen.Gen Double
genDouble lo hi =
  fmap (\n -> lo + (hi - lo) * fromIntegral n / 1e6) $
    Gen.inRange $ Range.between (0, 1000000 :: Word)

--------------------------------------------------------------------------------
-- Polyline reference
--------------------------------------------------------------------------------

-- | Approximate the arc length of a parametric curve by a polyline with @n@
-- segments.
polylineArcLength :: Int -> (Double -> ℝ 2) -> Double
polylineArcLength n f =
  sum ( zipWith ( distance @( T (ℝ 2) ) ) pts ( tail pts ) )
  where
    pts = [ f ( fromIntegral i / fromIntegral n )
          | i <- [0 .. n]
          ]

-- | Reference arc-length implementation, using Romberg's method.
referenceArcLength :: ( Double -> ℝ 2 ) -> Double
referenceArcLength f = ( 4 * l2 - l1 ) / 3
  where
    n  = 50000 -- Number of samples used for reference implementation.
    l1 = polylineArcLength       n   f
    l2 = polylineArcLength ( 2 * n ) f

--------------------------------------------------------------------------------
-- Error computation
--------------------------------------------------------------------------------

data Result
  = Result
  { referenceResult     :: !Double
  , methodResult        :: !Double
  , methodRelativeError :: !Double
  }
  deriving stock ( Show, Generic )
  deriving anyclass NFData

-- | Compute (impl, ref, relErr) for the quadratic arc-length test.
--
-- Endpoints fixed at p0=(0,0) and p2=(1,0); control point p1=(cx,cy).
quadResult :: ArcLengthOptions -> ℝ 2 -> Result
quadResult opts p1 =
  Result
    { referenceResult     = ref
    , methodResult        = impl
    , methodRelativeError = relErr
    }
  where
    p0    = ℝ2 0 0
    p2    = ℝ2 1 0
    bez   = Quadratic.Bezier { p0, p1, p2 }
    curve = Bezier2To
      { controlPoint = p1
      , curveEnd     = NextPoint p2
      , curveData    = ()
      }
    impl   = totalArcLength
           $ curveArcLengthParametrisation @( T ( ℝ 2 ) ) opts p0 curve
    ref    = referenceArcLength $ Quadratic.bezier @( T ( ℝ 2 ) ) bez
    relErr = abs (impl - ref) / max ref 1e-15

-- | Compute (impl, ref, relErr) for the cubic arc-length test.
--
-- Fixed points p0=(0,0), p1=(1,0), p3=(1,1); free control point p2=(cx,cy).
cubicResult :: ArcLengthOptions -> ℝ 2 -> Result
cubicResult opts p2 =
  Result
    { referenceResult     = ref
    , methodResult        = impl
    , methodRelativeError = relErr
    }
  where
    p0    = ℝ2 0 0
    p1    = ℝ2 1 0
    p3    = ℝ2 1 1
    bez   = Cubic.Bezier { p0, p1, p2, p3 }
    curve = Bezier3To
      { controlPoint1 = p1
      , controlPoint2 = p2
      , curveEnd      = NextPoint p3
      , curveData     = ()
      }
    impl   = totalArcLength
           $ curveArcLengthParametrisation @( T ( ℝ 2 ) ) opts p0 curve
    ref    = referenceArcLength $ Cubic.bezier @( T ( ℝ 2 ) ) bez
    relErr = abs (impl - ref) / max ref 1e-15

--------------------------------------------------------------------------------
-- Properties
--------------------------------------------------------------------------------

-- | Relative-error tolerance for arc-length tests.
arcLengthTol :: Double
arcLengthTol = 1e-5

testOptions :: Falsify.TestOptions
testOptions = Falsify.TestOptions
  { Falsify.expectFailure      = Falsify.DontExpectFailure
  , Falsify.overrideVerbose    = Nothing
  , Falsify.overrideNumTests   = Just 400
  , Falsify.overrideMaxShrinks = Just 0 -- shrinking isn't really useful here
  , Falsify.overrideMaxRatio   = Nothing
  }

-- | Arc length of a quadratic Bézier with endpoints (0,0), (1,0) and a random
-- control point (cx, cy) in [-3, 3]².
prop_quadratic :: ArcLengthOptions -> Property ()
prop_quadratic opts = do
  cx <- Falsify.genWith (\v -> Just $ "cx = " ++ show v) $ genDouble -3 3
  cy <- Falsify.genWith (\v -> Just $ "cy = " ++ show v) $ genDouble -3 3

  let
    c = ℝ2 cx cy
    Result
      { referenceResult     = ref
      , methodResult        = impl
      , methodRelativeError = relErr
      } = quadResult opts c

  Falsify.info $
    "{ c = " ++ show c ++
    ", impl = " ++ show impl ++
    ", ref = "  ++ show ref  ++
    ", relErr = " ++ show relErr ++ " }"

  Falsify.assert $
    Falsify.Prop.relatedBy ("< tolerance", (<))
      .$ ("relErr",    relErr)
      .$ ("tolerance", arcLengthTol)

-- | Arc length of a cubic Bézier with fixed points p0=(0,0), p1=(1,0), p3=(1,1)
-- and a random free control point p2=(cx, cy) in [-3, 3]².
prop_cubic :: ArcLengthOptions -> Property ()
prop_cubic opts = do
  cx <- Falsify.genWith (\v -> Just $ "cx = " ++ show v) $ genDouble -3 3
  cy <- Falsify.genWith (\v -> Just $ "cy = " ++ show v) $ genDouble -3 3

  let
    c = ℝ2 cx cy
    Result
      { referenceResult     = ref
      , methodResult        = impl
      , methodRelativeError = relErr
      } = cubicResult opts c

  Falsify.info $
    "{ c = " ++ show c ++
    ", impl = " ++ show impl ++
    ", ref = "  ++ show ref  ++
    ", relErr = " ++ show relErr ++ " }"

  Falsify.assert $
    Falsify.Prop.relatedBy ("< tolerance", (<))
      .$ ("relErr",    relErr)
      .$ ("tolerance", arcLengthTol)

testGroup :: ArcLengthOptions -> Tasty.TestTree
testGroup opts =
  Tasty.testGroup ( integratorName opts )
    [ testPropertyWith testOptions "Quadratic Bezier" $ prop_quadratic opts
    , testPropertyWith testOptions "Cubic Bezier"     $ prop_cubic     opts
    ]

allTests :: [ Tasty.TestTree ]
allTests = map ( testGroup . snd ) integrators

--------------------------------------------------------------------------------
-- CSV output
--------------------------------------------------------------------------------

writeArcLengthData :: FilePath -> ArcLengthOptions -> Double -> IO ()
writeArcLengthData fp opts gridSize = do
  putStrLn $
    unlines
      [ "Computing arc-length data for integrator" ++ integratorName opts ++ "."
      , "Writing to: " ++ fp
      ]
  writeFile fp $ displayDataPoints dataPoints
  putStrLn $
    unlines
      [ "  Worst error, at " ++ show wp ++ ": " ++ show werr
      , "   Mean error: " ++ show avg_err
      ]

  where
    chunkSize = 2500 -- one row per spark

    dataPoints :: [ ( ℝ 2, Result ) ]
    dataPoints =
      withStrategy ( parListChunk chunkSize rdeepseq )
        [ ( p, res )
        | x <- [ -2, -2 + gridSize .. 3 ]
        , y <- [ -2, -2 + gridSize .. 3 ]
        , let p = ℝ2 x y
              res = cubicResult opts p
        ]

    getErr :: ( ℝ 2, Result ) -> Double
    getErr ( _xy, res ) = methodRelativeError res

    ( wp, werr ) = maximumBy ( comparing getErr ) dataPoints

    log_avgerr = sum (map (logBase 10 . getErr) dataPoints) / fromIntegral ( length dataPoints )
    avg_err = 10 ** log_avgerr

displayDataPoints
  :: [ ( ℝ 2, Result ) ]
  -> String
displayDataPoints pts =
  intercalate "\n" ( map displayDataPoint pts ) ++ "\n"

displayDataPoint
  :: ( ℝ 2, Result )
  -> String
displayDataPoint ( ℝ2 x y, res ) =
  intercalate ", " $ map displayDouble [ x, y, methodRelativeError res ]

displayDouble :: Double -> String
displayDouble d = show d

--------------------------------------------------------------------------------
-- Command-line interface
--------------------------------------------------------------------------------

data Command
  = RunTests
  | WriteCSV
      { csvOutputPath :: !FilePath
      , csvIntegrator :: !ArcLengthOptions
      , csvGridSize   :: !Double
      }

commandParser :: OptParse.Parser Command
commandParser =
  OptParse.subparser
    (  OptParse.command "test"
         ( OptParse.info ( pure RunTests )
             ( OptParse.progDesc "(default) Run arc-length integrator testsuite" ) )
    <> OptParse.command "csv"
         ( OptParse.info ( csvParser OptParse.<**> OptParse.helper )
             ( OptParse.progDesc
                 "Write arc-length error data to a CSV file" ) )
    )
  -- No command: run the tests.
  <|> pure RunTests
  where
    csvParser :: OptParse.Parser Command
    csvParser =
      WriteCSV
        <$> OptParse.option OptParse.str
              ( OptParse.long "output"
             <> OptParse.short 'o'
             <> OptParse.metavar "FILE"
             <> OptParse.help "Path of CSV output file" )
        <*> OptParse.argument ( OptParse.eitherReader parseIntegrator )
              ( OptParse.metavar "INTEGRATOR"
             <> OptParse.help ( "Integration method to use: "
                    ++ intercalate ", " ( map ( NE.head . fst ) integrators ) ) )
        <*> OptParse.option OptParse.auto
              ( OptParse.long "grid-size"
             <> OptParse.short 'g'
             <> OptParse.metavar "STEP"
             <> OptParse.value 5e-2
             <> OptParse.showDefault
             <> OptParse.help "Grid step size" )

optParseOpts :: OptParse.ParserInfo Command
optParseOpts = OptParse.info ( commandParser OptParse.<**> OptParse.helper )
  (  OptParse.fullDesc
  <> OptParse.header "test-arc-length – test-suite for arc-length integrators"
  )

integrators :: [ ( NE.NonEmpty String, ArcLengthOptions ) ]
integrators =
  map ( \ i -> ( integratorNames i, i ) )
    [ gravesenOpts
    , gaussLegendreOpts
    , clenshawCurtisOpts
    , tanhSinhOpts
    ]

  where
    gravesenOpts =
      defaultArcLengthOptions
        { arcLengthIntegrator = Gravesen { gravesenTol = 1e-5 }
        }
    gaussLegendreOpts =
      defaultArcLengthOptions
        { arcLengthIntegrator   = GaussLegendre
            { gaussLegendreDegree = 20
            , regulariser         = Just defaultRegulariser
            }
        , splitAtCriticalPoints = True
        }
    clenshawCurtisOpts =
      defaultArcLengthOptions
        { arcLengthIntegrator   = ClenshawCurtis
            { clenshawCurtisDegree = 20
            , regulariser          = Just defaultRegulariser
            }
        , splitAtCriticalPoints = True
        }
    tanhSinhOpts =
      defaultArcLengthOptions
        { arcLengthIntegrator  = TanhSinh
            { tanhSinhNodes    = 32
            , tanhSinhStepSize = 0.1
            }
        , splitAtCriticalPoints = True
        }

integratorName :: ArcLengthOptions -> String
integratorName = NE.head . integratorNames

integratorNames :: ArcLengthOptions -> NE.NonEmpty String
integratorNames opts =
  case arcLengthIntegrator opts of
    Gravesen {}       -> "Gravesen" NE.:| [ "AdaptiveGravesen" ]
    GaussLegendre {}  -> "Gauss–Legendre" NE.:| [ "Gauss-Legendre", "GaussLegendre", "Gauss", "Legendre" ]
    ClenshawCurtis {} -> "Clenshaw–Curtis" NE.:| [  "Clenshaw-Curtis", "ClenshawCurtis", "Clenshaw", "Curtis" ]
    TanhSinh {}       -> "TanhSinh" NE.:| [ "Tanh–Sinh", "Tanh-Sinh", "Tanh", "Sinh" ]

parseIntegrator :: String -> Either String ArcLengthOptions
parseIntegrator intStr =
  case mapMaybe parseOne integrators of
    [ opts ] -> Right opts
    [ ] ->
      Left $ "Unknown integrator " ++ show intStr
          ++ ".\n  Available integrators: "
          ++ intercalate ", " integratorNms ++ "."
    _ ->
      Left $ "Ambiguous integrator " ++ show intStr
          ++ ".\n  Available integrators: "
          ++ intercalate ", " integratorNms ++ "."
  where
    integratorNms = map ( NE.head . fst ) integrators
    parseOne :: ( NE.NonEmpty String, ArcLengthOptions ) -> Maybe ArcLengthOptions
    parseOne ( nms, integrator ) = do
      guard $ any ( ( map toLower intStr == ) . map toLower ) nms
      return integrator

--------------------------------------------------------------------------------
-- Entry point
--------------------------------------------------------------------------------

main :: IO ()
main = withCP65001 $ do
  cmd <- OptParse.execParser optParseOpts
  case cmd of
    RunTests ->
      withArgs [] $
        Tasty.defaultMain $
          Tasty.testGroup "Arc-length integrators" allTests
    WriteCSV { csvOutputPath, csvIntegrator, csvGridSize } ->
      writeArcLengthData csvOutputPath csvIntegrator csvGridSize