feat(plot): include support for indicators

rodrigo-brito · rodrigo-brito · commit 0b69e8ea96e9 · 2021-10-03T01:10:08.000-03:00
diff --git a/examples/backtesting/main.go b/examples/backtesting/main.go
@@ -56,7 +56,9 @@ func main() {
 
 	chart := plot.NewChart(plot.WithIndicators(
 		plot.EMA(9, "red"),
-		plot.RSI(14, "purple"),
+		plot.EMA(80, "orange"),
+		plot.RSI(14, "green"),
+		plot.Stoch(8, 3, "red", "blue"),
 	))
 
 	bot, err := ninjabot.NewBot(
diff --git a/plot/assets/chart.html b/plot/assets/chart.html
@@ -12,6 +12,11 @@
   </head>
   <script defer src="/assets/chart.js"></script>
   <style>
+    body {
+      margin: 0;
+      padding: 0;
+    }
+
     .coin {
       background: rgb(68 122 219);
       padding: 10px;
@@ -41,11 +46,13 @@
     }
   </style>
   <body>
-    <ul>
-      {{range $val := .pairs}}
+    <nav class="menu">
+      <ul>
+        {{range $val := .pairs}}
         <li><a class="coin" href="/?pair={{ $val }}">{{ $val }}</a></li>
-      {{end}}
-    </ul>
+        {{end}}
+      </ul>
+    </nav>
     <div id="graph"></div>
   </body>
 </html>
diff --git a/plot/assets/chart.js b/plot/assets/chart.js
@@ -18,8 +18,8 @@ document.addEventListener("DOMContentLoaded", function () {
             low: unpack(data.candles, "low"),
             high: unpack(data.candles, "high"),
             type: "candlestick",
-            xaxis: "x",
-            yaxis: "y",
+            xaxis: "x1",
+            yaxis: "y1",
         };
 
         const points = [];
@@ -42,6 +42,8 @@ document.addEventListener("DOMContentLoaded", function () {
                     y: candle.low,
                     xref: "x",
                     yref: "y",
+                    xaxis: "x1",
+                    yaxis: "y1",
                     text: "B",
                     hovertext: `${order.time}
                         <br>ID: ${order.id}
@@ -70,6 +72,7 @@ document.addEventListener("DOMContentLoaded", function () {
                     annotation.ay = -20;
                     annotation.valign = "top";
                 }
+
                 annotations.push(annotation);
             });
         });
@@ -80,6 +83,8 @@ document.addEventListener("DOMContentLoaded", function () {
             name: "Buy Points",
             x: unpack(buyPoints, "time"),
             y: unpack(buyPoints, "position"),
+            xaxis: "x1",
+            yaxis: "y1",
             mode: 'markers',
             type: 'scatter',
             marker: {
@@ -90,31 +95,81 @@ document.addEventListener("DOMContentLoaded", function () {
             name: "Sell Points",
             x: unpack(sellPoints, "time"),
             y: unpack(sellPoints, "position"),
+            xaxis: "x1",
+            yaxis: "y1",
             mode: 'markers',
             type: 'scatter',
             marker: {
                 color: "red",
             }
         };
 
-        var layout = {
-            dragmode: "pan",
+        const standaloneIndicators = data.indicators.reduce((total, indicator) => {
+            if (!indicator.overlay) {
+                return total + 1;
+            }
+            return total;
+        }, 0);
+
+        let layout = {
+            template: "ggplot2",
+            dragmode: "zoom",
             margin: {
-                r: 10,
                 t: 25,
-                b: 40,
-                l: 60,
             },
             showlegend: true,
             xaxis: {
-                autorange: true
+                autorange: true,
+                rangeslider: {visible: false},
+                showline: true,
+                anchor: standaloneIndicators > 0 ? "y2" : "y1"
             },
             yaxis: {
-                autorange: true
+                domain: standaloneIndicators > 0 ? [0.5, 1] : [0, 1],
+                autorange: true,
+                mirror: true,
+                showline: true,
+                gridcolor: "#ddd"
             },
+            hovermode: "x unified",
             annotations: annotations,
         };
 
-        Plotly.newPlot("graph", [candleStickData, buyData, sellData], layout);
+        let plotData = [candleStickData, buyData, sellData];
+        const indicatorsHeight = 0.49/standaloneIndicators;
+        let standaloneIndicatorIndex = 0;
+        data.indicators.forEach((indicator) => {
+            const axisNumber = standaloneIndicatorIndex+2;
+            if (!indicator.overlay) {
+                const heightStart = standaloneIndicatorIndex * indicatorsHeight;
+                layout["yaxis"+axisNumber] = {
+                    domain: [heightStart, heightStart + indicatorsHeight],
+                    autorange: true,
+                    mirror: true,
+                    showline: true,
+                    linecolor: "black",
+                    title: indicator.name
+                };
+                standaloneIndicatorIndex++;
+            }
+
+            indicator.metrics.forEach(metric => {
+                const data = {
+                    title: indicator.name,
+                    name: indicator.name + (metric.name && " - " + metric.name),
+                    x: metric.time,
+                    y: metric.value,
+                    type: metric.style,
+                    color: metric.color,
+                    xaxis: "x1",
+                    yaxis: "y1",
+                };
+                if (!indicator.overlay) {
+                    data.yaxis = "y"+axisNumber;
+                }
+                plotData.push(data);
+            })
+        });
+        Plotly.newPlot("graph", plotData, layout);
     })
 });
diff --git a/plot/indicator.go b/plot/indicator.go
@@ -47,14 +47,17 @@ func (e ema) Overlay() bool {
 }
 
 func (e *ema) Load(dataframe *model.Dataframe) {
-	e.Values = talib.Ema(dataframe.Close, e.Period)
-	e.Time = dataframe.Time
+	if len(dataframe.Time) < e.Period {
+		return
+	}
+
+	e.Values = talib.Ema(dataframe.Close, e.Period)[e.Period:]
+	e.Time = dataframe.Time[e.Period:]
 }
 
 func (e ema) Metrics() []Metric {
 	return []Metric{
 		{
-			Name:   "value",
 			Style:  "line",
 			Color:  e.Color,
 			Values: e.Values,
@@ -64,7 +67,7 @@ func (e ema) Metrics() []Metric {
 }
 
 func RSI(period int, color string) Indicator {
-	return &ema{
+	return &rsi{
 		Period: period,
 		Color:  color,
 	}
@@ -86,18 +89,78 @@ func (e rsi) Overlay() bool {
 }
 
 func (e *rsi) Load(dataframe *model.Dataframe) {
-	e.Values = talib.Rsi(dataframe.Close, e.Period)
-	e.Time = dataframe.Time
+	if len(dataframe.Time) < e.Period {
+		return
+	}
+
+	e.Values = talib.Rsi(dataframe.Close, e.Period)[e.Period:]
+	e.Time = dataframe.Time[e.Period:]
 }
 
 func (e rsi) Metrics() []Metric {
 	return []Metric{
 		{
-			Name:   "value",
 			Color:  e.Color,
 			Style:  "line",
 			Values: e.Values,
 			Time:   e.Time,
 		},
 	}
 }
+
+func Stoch(k, d int, colork, colord string) Indicator {
+	return &stoch{
+		PeriodK: k,
+		PeriodD: d,
+		ColorK:  colork,
+		ColorD:  colord,
+	}
+}
+
+type stoch struct {
+	PeriodK int
+	PeriodD int
+	ColorK  string
+	ColorD  string
+	ValuesK model.Series
+	ValuesD model.Series
+	Time    []time.Time
+}
+
+func (e stoch) Name() string {
+	return fmt.Sprintf("STOCH(%d, %d)", e.PeriodK, e.PeriodD)
+}
+
+func (e stoch) Overlay() bool {
+	return false
+}
+
+func (e *stoch) Load(dataframe *model.Dataframe) {
+	if len(dataframe.Time) < e.PeriodK+e.PeriodD {
+		return
+	}
+
+	e.ValuesK, e.ValuesD = talib.Stoch(dataframe.High, dataframe.Low, dataframe.Close, e.PeriodK, e.PeriodD, talib.SMA, e.PeriodD, talib.SMA)
+	e.ValuesK = e.ValuesK[e.PeriodK+e.PeriodD:]
+	e.ValuesD = e.ValuesD[e.PeriodK+e.PeriodD:]
+	e.Time = dataframe.Time[e.PeriodK+e.PeriodD:]
+}
+
+func (e stoch) Metrics() []Metric {
+	return []Metric{
+		{
+			Color:  e.ColorK,
+			Name:   "K",
+			Style:  "line",
+			Values: e.ValuesK,
+			Time:   e.Time,
+		},
+		{
+			Color:  e.ColorD,
+			Name:   "D",
+			Style:  "line",
+			Values: e.ValuesD,
+			Time:   e.Time,
+		},
+	}
+}
