Changed returns_histogram and monthly_returns_heatmap to work with bt.balance

backtester/statistics/__init__.py

@@ -0,0 +1 @@
+from .charts import monthly_returns_heatmap

backtester/statistics/charts.py

@@ -8,23 +8,13 @@ def returns_chart(report):
     time_interval = alt.selection(type="interval", encodings=["x"])
 
     # Area plot
-    areas = alt.Chart().mark_area(opacity=0.6).encode(
-        x=alt.X(
-            "index:T",
-            axis=alt.Axis(title="Date"),
-            scale={"domain": time_interval.ref()}),
-        y=alt.Y("% Price:Q", axis=alt.Axis(format="%")))
+    areas = alt.Chart().mark_area(opacity=0.6).encode(x='index:T',
+                                                      y=alt.Y("accumulated return:Q", axis=alt.Axis(format="%")))
 
     # Nearest point selector
-    nearest = alt.selection(
-        type="single",
-        nearest=True,
-        on="mouseover",
-        fields=["index"],
-        empty="none")
+    nearest = alt.selection(type="single", nearest=True, on="mouseover", fields=["index"], empty="none")
 
-    points = areas.mark_point().encode(
-        opacity=alt.condition(nearest, alt.value(1), alt.value(0)))
+    points = areas.mark_point().encode(opacity=alt.condition(nearest, alt.value(1), alt.value(0)))
 
     # Transparent date selector
     selectors = alt.Chart().mark_point().encode(
@@ -32,18 +22,17 @@ def returns_chart(report):
         opacity=alt.value(0),
     ).add_selection(nearest)
 
-    text = areas.mark_text(
-        align="left", dx=5, dy=-5).encode(
-            text=alt.condition(nearest, "% Price:Q", alt.value(" ")))
+    text = areas.mark_text(align="left", dx=5,
+                           dy=-5).encode(text=alt.condition(nearest, "accumulated return:Q", alt.value(" ")))
 
-    layered = alt.layer(
-        areas,
-        selectors,
-        points,
-        text,
-        width=700,
-        height=350,
-        title="Returns over time")
+    layered = alt.layer(areas.encode(
+        alt.X('index:T', axis=alt.Axis(title="date"), scale=alt.Scale(domain=time_interval))),
+                        selectors,
+                        points,
+                        text,
+                        width=700,
+                        height=350,
+                        title="Returns over time")
 
     lower = areas.properties(width=700, height=70).add_selection(time_interval)
 
@@ -53,28 +42,22 @@ def returns_chart(report):
 
 
 def returns_histogram(report):
-    bar = alt.Chart(report).mark_bar().encode(
-        x=alt.X(
-            "Interval Change:Q",
-            bin=alt.BinParams(maxbins=100),
-            axis=alt.Axis(format='%')),
-        y="count():Q")
+    bar = alt.Chart(report).mark_bar().encode(x=alt.X("% change:Q",
+                                                      bin=alt.BinParams(maxbins=100),
+                                                      axis=alt.Axis(format='%')),
+                                              y="count():Q")
     return bar
 
 
 def monthly_returns_heatmap(report):
-    monthly_returns = report.resample(
-        "M")["Total Portfolio"].last().pct_change().reset_index()
-    monthly_returns.columns = ["Date", "Monthly Returns"]
+    resample = report.resample("M")["capital"].last()
+    monthly_returns = resample.pct_change().reset_index()
+    monthly_returns['capital'].iat[0] = resample.iloc[0] / report.iloc[0]['capital'] - 1
+    monthly_returns.columns = ["date", "capital"]
 
     chart = alt.Chart(monthly_returns).mark_rect().encode(
-        alt.X("year(Date):O", title="Year"),
-        alt.Y("month(Date):O", title="Month"),
-        alt.Color(
-            "mean(Monthly Returns)",
-            title="Return",
-            scale=alt.Scale(scheme="redyellowgreen")),
-        alt.Tooltip("mean(Monthly Returns)",
-                    format=".2f")).properties(title="Average Monthly Returns")
+        alt.X("year(date):O", title="Year"), alt.Y("month(date):O", title="Month"),
+        alt.Color("mean(capital)", title="Return", scale=alt.Scale(scheme="redyellowgreen")),
+        alt.Tooltip("mean(capital)", format=".2f")).properties(title="Monthly Returns")
 
     return chart