BarAggregator: support for bars

src/ATrade/BarAggregator.hs

@@ -18,6 +18,7 @@ module ATrade.BarAggregator (
   BarAggregator(..),
   mkAggregatorFromBars,
   handleTick,
+  handleBar,
   hmsToDiffTime
 ) where
 
@@ -133,3 +134,53 @@ handleTick tick = runState $ do
           barLow = barClose bar,
           barClose = barClose bar,
           barVolume = 0 }
+
+handleBar :: Bar -> BarAggregator -> (Maybe Bar, BarAggregator)
+handleBar bar = runState $ do
+  tws <- gets tickTimeWindows
+  mybars <- gets bars
+  if (any (isInTimeInterval bar) tws)
+    then
+      case M.lookup (barSecurity bar) mybars of
+        Just series -> case bsBars series of
+          (b:bs) -> do
+            let currentBn = barNumber (barTimestamp b) (tfSeconds $ bsTimeframe series)
+            if currentBn == barNumber (barTimestamp bar) (tfSeconds $ bsTimeframe series)
+              then do
+                lBars %= M.insert (barSecurity bar) series { bsBars = updateBar b bar : bs }
+                return Nothing
+              else
+                if barEndTime b (tfSeconds $ bsTimeframe series) == barTimestamp bar
+                  then do
+                    lBars %= M.insert (barSecurity bar) series { bsBars = emptyBarFrom bar : (updateBar b bar : bs) }
+                    return . Just $ updateBar b bar
+                  else do
+                    lBars %= M.insert (barSecurity bar) series { bsBars = bar : b : bs }
+                    return . Just $ b
+          _      -> do
+            lBars %= M.insert (barSecurity bar) series { bsBars = [bar] }
+            return Nothing
+        _ -> return Nothing
+    else
+      return Nothing
+  where
+    isInTimeInterval bar' (a, b) = (utctDayTime . barTimestamp) bar' >= a && (utctDayTime . barTimestamp) bar' <= b
+    updateBar !bar' newbar =
+      let newHigh = max (barHigh bar') (barHigh newbar)
+          newLow = min (barLow bar') (barLow newbar) in
+        bar' {
+            barTimestamp = barTimestamp newbar,
+            barHigh = newHigh,
+            barLow = newLow,
+            barClose = barClose newbar,
+            barVolume = barVolume bar' + (abs . barVolume $ newbar) }
+
+    emptyBarFrom bar' = Bar {
+          barSecurity = barSecurity bar',
+          barTimestamp = barTimestamp bar',
+          barOpen = barClose bar',
+          barHigh = barClose bar',
+          barLow = barClose bar',
+          barClose = barClose bar',
+          barVolume = 0 }
+

test/Test/BarAggregator.hs

@@ -23,6 +23,10 @@ unitTests = testGroup "BarAggregator" [
   , testOneTick
   , testTwoTicksInSameBar
   , testTwoTicksInDifferentBars
+  , testOneBar
+  , testTwoBarsInSameBar
+  , testTwoBarsInSameBarLastBar
+  , testNextBarAfterBarClose
   ]
 
 testUnknownBarSeries :: TestTree
@@ -93,3 +97,87 @@ testTwoTicksInDifferentBars = testCase "Two ticks - different bar" $ do
             timestamp = ts,
             value = fromDouble val,
             volume = 1 }
+
+testOneBar :: TestTree
+testOneBar  = testCase "One bar" $ do
+  let series = BarSeries "TEST_TICKER" (Timeframe 3600) []
+  let agg = mkAggregatorFromBars (M.fromList [("TEST_TICKER", series)]) [(0, 86400)]
+  let (mbar, newagg) = handleBar bar agg
+  mbar @?= Nothing
+  (bsBars <$> (M.lookup "TEST_TICKER" $ bars newagg)) @?= Just [Bar "TEST_TICKER" testTimestamp 12.00 18.00 10.00 12.00 68]
+  where
+    testTimestamp = (UTCTime (fromGregorian 1970 1 1) 60)
+    bar = Bar {
+            barSecurity = "TEST_TICKER",
+            barTimestamp = testTimestamp,
+            barOpen = fromDouble 12.00,
+            barHigh = fromDouble 18.00,
+            barLow = fromDouble 10.00,
+            barClose = fromDouble 12.00,
+            barVolume = 68 }
+
+
+testTwoBarsInSameBar :: TestTree
+testTwoBarsInSameBar = testCase "Two bars (smaller timeframe) - same bar" $ do
+  let series = BarSeries "TEST_TICKER" (Timeframe 600) []
+  let agg = mkAggregatorFromBars (M.fromList [("TEST_TICKER", series)]) [(0, 86400)]
+  let (mbar, newagg) = handleBar (bar testTimestamp1 12.00 13.00 10.00 11.00 1) agg
+  mbar @?= Nothing
+  let (mbar', newagg') = handleBar (bar testTimestamp2 12.00 15.00 11.00 12.00 2) newagg
+  mbar' @?= Nothing
+  (bsBars <$> (M.lookup "TEST_TICKER" $ bars newagg')) @?= Just [Bar "TEST_TICKER" testTimestamp2 12.00 15.00 10.00 12.00 3]
+  where
+    testTimestamp1 = (UTCTime (fromGregorian 1970 1 1) 60)
+    testTimestamp2 = (UTCTime (fromGregorian 1970 1 1) 120)
+    bar ts o h l c v = Bar {
+            barSecurity = "TEST_TICKER",
+            barTimestamp = ts,
+            barOpen = fromDouble o,
+            barHigh = fromDouble h,
+            barLow = fromDouble l,
+            barClose = fromDouble c,
+            barVolume = v }
+
+testTwoBarsInSameBarLastBar :: TestTree
+testTwoBarsInSameBarLastBar = testCase "Two bars (smaller timeframe) - same bar: last bar is exactly at the end of the bigger tf bar" $ do
+  let series = BarSeries "TEST_TICKER" (Timeframe 600) []
+  let agg = mkAggregatorFromBars (M.fromList [("TEST_TICKER", series)]) [(0, 86400)]
+  let (mbar, newagg) = handleBar (bar testTimestamp1 12.00 13.00 10.00 11.00 1) agg
+  mbar @?= Nothing
+  let (mbar', newagg') = handleBar (bar testTimestamp2 12.00 15.00 11.00 12.00 2) newagg
+  let expectedBar = Bar "TEST_TICKER" testTimestamp2 12.00 15.00 10.00 12.00 3
+  mbar' @?= Just expectedBar
+  (head . tail <$> bsBars <$> (M.lookup "TEST_TICKER" $ bars newagg')) @?= Just expectedBar
+  where
+    testTimestamp1 = (UTCTime (fromGregorian 1970 1 1) 560)
+    testTimestamp2 = (UTCTime (fromGregorian 1970 1 1) 600)
+    bar ts o h l c v = Bar {
+            barSecurity = "TEST_TICKER",
+            barTimestamp = ts,
+            barOpen = fromDouble o,
+            barHigh = fromDouble h,
+            barLow = fromDouble l,
+            barClose = fromDouble c,
+            barVolume = v }
+
+testNextBarAfterBarClose :: TestTree
+testNextBarAfterBarClose = testCase "Three bars (smaller timeframe) - next bar after bigger tf bar close" $ do
+  let series = BarSeries "TEST_TICKER" (Timeframe 600) []
+  let agg = mkAggregatorFromBars (M.fromList [("TEST_TICKER", series)]) [(0, 86400)]
+  let (_, newagg) = handleBar (bar testTimestamp1 12.00 13.00 10.00 11.00 1) agg
+  let (_, newagg') = handleBar (bar testTimestamp2 12.00 15.00 11.00 12.00 2) newagg
+  let (_, newagg'') = handleBar (bar testTimestamp3 12.00 15.00 11.00 12.00 12) newagg'
+  let expectedBar = Bar "TEST_TICKER" testTimestamp3 12.00 15.00 11.00 12.00 12
+  (head <$> bsBars <$> (M.lookup "TEST_TICKER" $ bars newagg'')) @?= Just expectedBar
+  where
+    testTimestamp1 = (UTCTime (fromGregorian 1970 1 1) 560)
+    testTimestamp2 = (UTCTime (fromGregorian 1970 1 1) 600)
+    testTimestamp3 = (UTCTime (fromGregorian 1970 1 1) 660)
+    bar ts o h l c v = Bar {
+            barSecurity = "TEST_TICKER",
+            barTimestamp = ts,
+            barOpen = fromDouble o,
+            barHigh = fromDouble h,
+            barLow = fromDouble l,
+            barClose = fromDouble c,
+            barVolume = v }