""" Implementation of a standard financial plot visualization using Chaco renderers and scales. """ # Major library imports from numpy import abs, arange, cumprod, linspace, random, argmin, choose, vstack import time # Enthought library imports from enthought.enable.api import Component, ComponentEditor from enthought.traits.api import HasTraits, Instance, Enum from enthought.traits.ui.api import View, VGroup, Item, EnumEditor # Chaco imports from enthought.chaco.api import ArrayDataSource, DataRange1D, \ LinearMapper, PlotAxis, add_default_grids, OverlayPlotContainer from enthought.chaco.tools.api import PanTool, ZoomTool from enthought.chaco.scales.api import CalendarScaleSystem from enthought.chaco.scales_tick_generator import ScalesTickGenerator try: from enthought.chaco.hloc_renderer import OHLCPlot, PriceDataSource, \ CandleOHLCPlot, HeikinAshiPlot except ImportError: try: from ohlc_renderer import OHLCPlot, PriceDataSource, \ CandleOHLCPlot, HeikinAshiPlot except ImportError: raise ImportError, "Missing HLOC renderer module, unable to run this demo" def create_dates(numpoints, units="days"): """ Returns **numpoints** number of dates that evenly bracket the current date and time. **units** should be one of "weeks", "days", "hours" "minutes", or "seconds". """ units_map = { "weeks" : 7*24*3600, "days" : 24*3600, "hours" : 3600, "minutes" : 60, "seconds" : 1 } now = time.time() dt = units_map[units] dates = linspace(now, now+numpoints*dt, numpoints) return dates class MyApp(HasTraits): container = Instance(Component) plottype = Enum("OHLC", "Candlestick", "Heikin Ashi") view = View(VGroup( Item("plottype", show_label=False, width=-250), Item("container", editor=ComponentEditor(size=(800,600)), show_label=False) ), resizable = True, title = "OHLC Plot") _renderer_map = {"OHLC": OHLCPlot, "Candlestick": CandleOHLCPlot, "Heikin Ashi": HeikinAshiPlot} def _plottype_changed(self): oldplot = self.plot newplot = self._create_plot(oldplot.index, oldplot.prices) newplot.index_range.set_bounds(oldplot.index_range.low, oldplot.index_range.high) self.container.remove(oldplot) self.container.add(newplot) self.container.request_redraw() self.plot = newplot def _create_plot(self, times, prices): cls = self._renderer_map[self.plottype] if self.plot is None: index_mapper = LinearMapper(range=DataRange1D(times)) value_mapper = LinearMapper(range=DataRange1D(prices)) else: index_mapper = self.plot.index_mapper value_mapper = self.plot.value_mapper price_plot = cls(times = times, prices = prices, index_mapper = index_mapper, value_mapper = value_mapper, bgcolor = "white", border_visible = True) # Set the plot's bottom axis to use the Scales ticking system ticker = ScalesTickGenerator(scale=CalendarScaleSystem()) bottom_axis = PlotAxis(price_plot, orientation="bottom", tick_generator = ticker) price_plot.overlays.append(bottom_axis) price_plot.overlays.append(PlotAxis(price_plot, orientation="left")) hgrid, vgrid = add_default_grids(price_plot) vgrid.tick_generator = bottom_axis.tick_generator # Add pan and zoom price_plot.tools.append(PanTool(price_plot, constrain=True, constrain_direction="x")) price_plot.overlays.append( ZoomTool(price_plot, drag_button="right", always_on=True, tool_mode="range", axis="index")) return price_plot def _container_default(self): self.plot = None # Create the data and datasource objects # In order for the date axis to work, the index data points need to # be in units of seconds since the epoch. This is because we are using # the CalendarScaleSystem, whose formatters interpret the numerical values # as seconds since the epoch. numpoints = 500 index = create_dates(numpoints) returns = random.lognormal(0.00, 0.04, size=numpoints) average = 100.0 * cumprod(returns) high = average + abs(random.normal(0, 20.0, size=numpoints)) low = average - abs(random.normal(0, 20.0, size=numpoints)) delta = high - low open = low + delta * random.uniform(0.05, 0.95, size=numpoints) close = low + delta * random.uniform(0.05, 0.95, size=numpoints) price = vstack((open, high, low, close, average)) time_ds = ArrayDataSource(index) price_ds = PriceDataSource(price, sort_order="none") # Create the price plot price_plot = self._create_plot(time_ds, price_ds) self.plot = price_plot container = OverlayPlotContainer(padding=35) container.add(price_plot) return container demo = MyApp() if __name__ == "__main__": demo.configure_traits()