-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathplot_candles.py
105 lines (94 loc) · 4.13 KB
/
plot_candles.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
plt.rcParams["figure.figsize"] = (20,10)
def plot_candles(start_time, end_time, pricing, title=None,
volume_bars=False,
color_function=None,
overlays=None,
technicals=None,
technicals_titles=None):
""" Plots a candlestick chart using quantopian pricing data.
Author: Daniel Treiman
Args:
pricing: A pandas dataframe with columns ['open_price', 'close_price', 'high', 'low', 'volume']
title: An optional title for the chart
volume_bars: If True, plots volume bars
color_function: A function which, given a row index and price series, returns a candle color.
overlays: A list of additional data series to overlay on top of pricing. Must be the same length as pricing.
technicals: A list of additional data series to display as subplots.
technicals_titles: A list of titles to display for each technical indicator.
"""
pricing = pricing[start_time:end_time]
if overlays is not None:
overlays = [o[start_time:end_time] for o in overlays]
if technicals is not None:
technicals = [t[start_time:end_time] for t in technicals]
def default_color(index, open_price, close_price, low, high):
return 'g' if open_price[index] > close_price[index] else 'r'
color_function = color_function or default_color
overlays = overlays or []
technicals = technicals or []
technicals_titles = technicals_titles or []
open_price = pricing['open']
close_price = pricing['close']
low = pricing['low']
high = pricing['high']
oc_min = pd.concat([open_price, close_price], axis=1).min(axis=1)
oc_max = pd.concat([open_price, close_price], axis=1).max(axis=1)
subplot_count = 1
if volume_bars:
subplot_count = 2
if technicals:
subplot_count += len(technicals)
if subplot_count == 1:
fig, ax1 = plt.subplots(1, 1)
else:
ratios = np.insert(np.full(subplot_count - 1, 1), 0, 3)
fig, subplots = plt.subplots(subplot_count, 1, sharex=True, gridspec_kw={'height_ratios': ratios})
ax1 = subplots[0]
if title:
ax1.set_title(title)
x = np.arange(len(pricing))
candle_colors = [color_function(i, open_price, close_price, low, high) for i in x]
candles = ax1.bar(x, oc_max-oc_min, bottom=oc_min, color=candle_colors, linewidth=0)
lines = ax1.vlines(x , low, high, color=candle_colors, linewidth=1)#+ 0.4
ax1.xaxis.grid(False)
ax1.xaxis.set_tick_params(which='major', length=3.0, direction='in', top='off')
# Assume minute frequency if first two bars are in the same day.
frequency = 'minute' if (pricing.index[1] - pricing.index[0]).days == 0 else 'day'
time_format = '%d-%m-%Y'
if frequency == 'minute':
time_format = '%H:%M'
# Set X axis tick labels.
ticks = [date.strftime(time_format) for date in pricing.index]
space = max(int(len(ticks) / 20), 1)
for i, t in enumerate(ticks):
ticks[i] = t if i%space == 0 or i == len(ticks) - 1 else ''
plt.xticks(x, ticks, rotation='vertical')
for overlay in overlays:
ax1.plot(x, overlay)
# Plot volume bars if needed
if volume_bars:
ax2 = subplots[1]
volume = pricing['volume']
volume_scale = None
scaled_volume = volume
if volume.max() > 1000000:
volume_scale = 'M'
scaled_volume = volume / 1000000
elif volume.max() > 1000:
volume_scale = 'K'
scaled_volume = volume / 1000
ax2.bar(x, scaled_volume, color=candle_colors)
volume_title = 'Volume'
if volume_scale:
volume_title = 'Volume (%s)' % volume_scale
ax2.set_title(volume_title)
ax2.xaxis.grid(False)
# Plot additional technical indicators
for (i, technical) in enumerate(technicals):
ax = subplots[i - len(technicals)] # Technical indicator plots are shown last
ax.plot(x, technical)
if i < len(technicals_titles):
ax.set_title(technicals_titles[i])