content: -- | Tests "Synapse.NN" module and its submodules.
  
  
  module NNTest
      ( tests
      )
      where
  
  
  import Synapse.Tensors
  import qualified Synapse.Tensors.Vec as V
  
  import Synapse.NN.Models
  import Synapse.NN.Layers
  import Synapse.NN.Optimizers
  import Synapse.NN.LearningRates
  import Synapse.NN.Losses
  import Synapse.NN.Batching
  import Synapse.NN.Training
  
  import System.Random
  
  import Graphics.EasyPlot
  
  import Test.HUnit
  
  
  testSin :: Test  -- -3 sin (x + 5)
  testSin = TestLabel "testSin" $ TestCase $ do
      let sinFn x = (-3.0) * sin (x + 5.0)
      let model = buildSequentialModel (InputSize 1) [ Layer . layerDense 1
                                                     , Layer . layerActivation (Activation cos)
                                                     , Layer . layerDense 1
                                                     ] :: SequentialModel Double
  
      let dataset = Dataset $ V.fromList $ [Sample (singleton x) (sinFn $ singleton x) | x <- [-pi, -pi+0.2 .. pi]]
      (trainedModel, _, losses, _) <- train model
                                         (SGD 0.2 False)
                                         (Hyperparameters 500 16 dataset (LearningRate $ const 0.01) (Loss mse) V.empty)
                                         emptyCallbacks
                                         (mkStdGen 1)
      _ <- plot (PNG "test/plots/sin.png")
                [ Data2D [Title "predicted sin", Style Lines, Color Red] [Range (-pi) pi] [(x, unSingleton $ forward (singleton x) trainedModel) | x <- [-pi, -pi+0.05 .. pi]]
                , Data2D [Title "true sin", Style Lines, Color Green] [Range (-pi) pi] [(x, sinFn x) | x <- [-pi, -pi+0.05 .. pi]]
                ]
  
      let unpackedLosses = unRecordedMetric (unsafeIndex losses 0)
      let lastLoss = unsafeIndex unpackedLosses (V.size unpackedLosses - 1)
  
      assertBool "trained well enough" (lastLoss < 0.01)
  
  testSqrt :: Test  -- sqrt(x)
  testSqrt = TestLabel "testSqrt" $ TestCase $ do
      let sqrtFn x = sqrt x
      let model = buildSequentialModel (InputSize 1) [ Layer . layerDense 1
                                                     , Layer . layerActivation (Activation tanh)
                                                     , Layer . layerDense 1
                                                     ] :: SequentialModel Double
  
      let dataset = Dataset $ V.fromList $ [Sample (singleton x) (sqrtFn $ singleton x) | x <- [0.0, 0.2 .. 4.0]]
      (trainedModel, _, losses, _) <- train model
                                         (SGD 0.2 True)
                                         (Hyperparameters 500 16 dataset (LearningRate $ const 0.01) (Loss mse) V.empty)
                                         emptyCallbacks
                                         (mkStdGen 1)
      _ <- plot (PNG "test/plots/sqrt.png")
                 [ Data2D [Title "predicted sqrt", Style Lines, Color Red] [Range 0.0 6.0] $ [(x, unSingleton $ forward (singleton x) trainedModel) | x <- [0.0, 0.05 .. 4.0]]
                 , Data2D [Title "true sqrt", Style Lines, Color Green] [Range 0.0 6.0] $ [(x, sqrtFn x) | x <- [0.0, 0.05 .. 4.0]]
                 ]
  
      let unpackedLosses = unRecordedMetric (unsafeIndex losses 0)
      let lastLoss = unsafeIndex unpackedLosses (V.size unpackedLosses - 1)
  
      assertBool "trained well enough" (lastLoss < 0.01)
  
  testTrigonometry :: Test  -- sin(2.0 * cos(x) + 3.0) + 2.5
  testTrigonometry = TestLabel "testTrigonometry" $ TestCase $ do
      let trigonometryFn x = sin (2.0 * cos x + 3.0) + 2.5
      let model = buildSequentialModel (InputSize 1) [ Layer . layerDense 1
                                                     , Layer . layerActivation (Activation sin)
                                                     , Layer . layerDense 1
                                                     , Layer . layerActivation (Activation sin)
                                                     , Layer . layerDense 1
                                                     ] :: SequentialModel Double
  
      let dataset = Dataset $ V.fromList $ [Sample (singleton x) (trigonometryFn $ singleton x) | x <- [-(2.0 * pi),((-(2.0 * pi)) + 0.1)..(2.0 * pi)]]
                                           
      (trainedModel, _, losses, _) <- train model
                                         (SGD 0.3 True)
                                         (Hyperparameters 1000 1 dataset (LearningRate $ const 0.001) (Loss mse) V.empty)
                                         emptyCallbacks
                                         (mkStdGen 1)
  
      _ <- plot (PNG "test/plots/trigonometry.png")
                 [ Data2D [Title "predicted trigonometry", Style Lines, Color Red] [Range ((-2.0) * pi) (2.0 * pi)] $ [(x, unSingleton $ forward (singleton x) trainedModel) | x <- [-(2.0 * pi),((-(2.0 * pi)) + 0.1)..(2.0 * pi)]]
                 , Data2D [Title "true trigonometry", Style Lines, Color Green] [Range ((-2.0) * pi) (2.0 * pi)] $ [(x, trigonometryFn x) | x <- [-(2.0 * pi),((-(2.0 * pi)) + 0.1)..(2.0 * pi)]]
                 ]
  
      let unpackedLosses = unRecordedMetric (unsafeIndex losses 0)
      let lastLoss = unsafeIndex unpackedLosses (V.size unpackedLosses - 1)
  
      assertBool "trained well enough" (lastLoss < 0.01)
  
  
  tests :: Test
  tests = TestLabel "NNTest" $ TestList
      [ testSin
      , testSqrt
      , testTrigonometry
      ]
  
  