I'd like to add an arrow to a line plot with matplotlib like in the plot below (drawn with pgfplots).
How can I do (position and direction of the arrow should be parameters ideally)?
Here is some code to experiment.
from matplotlib import pyplot
import numpy as np
t = np.linspace(-2, 2, 100)
plt.plot(t, np.sin(t))
plt.show()
Thanks.
In my experience this works best by using annotate. Thereby you avoid the weird warping you get with ax.arrow which is somehow hard to control.
EDIT: I've wrapped it into a little function.
from matplotlib import pyplot as plt
import numpy as np
def add_arrow(line, position=None, direction='right', size=15, color=None):
"""
add an arrow to a line.
line: Line2D object
position: x-position of the arrow. If None, mean of xdata is taken
direction: 'left' or 'right'
size: size of the arrow in fontsize points
color: if None, line color is taken.
"""
if color is None:
color = line.get_color()
xdata = line.get_xdata()
ydata = line.get_ydata()
if position is None:
position = xdata.mean()
# find closest index
start_ind = np.argmin(np.absolute(xdata - position))
if direction == 'right':
end_ind = start_ind + 1
else:
end_ind = start_ind - 1
line.axes.annotate('',
xytext=(xdata[start_ind], ydata[start_ind]),
xy=(xdata[end_ind], ydata[end_ind]),
arrowprops=dict(arrowstyle="->", color=color),
size=size
)
t = np.linspace(-2, 2, 100)
y = np.sin(t)
# return the handle of the line
line = plt.plot(t, y)[0]
add_arrow(line)
plt.show()
It's not very intuitive but it works. You can then fiddle with the arrowprops dictionary until it looks right.
Just add a plt.arrow():
from matplotlib import pyplot as plt
import numpy as np
# your function
def f(t): return np.sin(t)
t = np.linspace(-2, 2, 100)
plt.plot(t, f(t))
plt.arrow(0, f(0), 0.01, f(0.01)-f(0), shape='full', lw=0, length_includes_head=True, head_width=.05)
plt.show()
EDIT: Changed parameters of arrow to include position & direction of function to draw.
Not the nicest solution, but should work:
import matplotlib.pyplot as plt
import numpy as np
def makeArrow(ax,pos,function,direction):
delta = 0.0001 if direction >= 0 else -0.0001
ax.arrow(pos,function(pos),pos+delta,function(pos+delta),head_width=0.05,head_length=0.1)
fun = np.sin
t = np.linspace(-2, 2, 100)
ax = plt.axes()
ax.plot(t, fun(t))
makeArrow(ax,0,fun,+1)
plt.show()
I know this doesn't exactly answer the question as asked, but I thought this could be useful to other people landing here. I wanted to include the arrow in my plot's legend, but the solutions here don't mention how. There may be an easier way to do this, but here is my solution:
To include the arrow in your legend, you need to make a custom patch handler and use the matplotlib.patches.FancyArrow object. Here is a minimal working solution. This solution piggybacks off of the existing solutions in this thread.
First, the imports...
import matplotlib.pyplot as plt
from matplotlib.legend_handler import HandlerPatch
import matplotlib.patches as patches
from matplotlib.lines import Line2D
import numpy as np
Now, we make a custom legend handler. This handler can create legend artists for any line-patch combination, granted that the line has no markers.
class HandlerLinePatch(HandlerPatch):
def __init__(self, linehandle=None, **kw):
HandlerPatch.__init__(self, **kw)
self.linehandle=linehandle
def create_artists(self, legend, orig_handle,
xdescent, ydescent, width,
height, fontsize, trans):
p = super().create_artists(legend, orig_handle,
xdescent, descent,
width, height, fontsize,
trans)
line = Line2D([0,width],[height/2.,height/2.])
if self.linehandle is None:
line.set_linestyle('-')
line._color = orig_handle._edgecolor
else:
self.update_prop(line, self.linehandle, legend)
line.set_drawstyle('default')
line.set_marker('')
line.set_transform(trans)
return [p[0],line]
Next, we write a function that specifies the type of patch we want to include in the legend - an arrow in our case. This is courtesy of Javier's answer here.
def make_legend_arrow(legend, orig_handle,
xdescent, ydescent,
width, height, fontsize):
p = patches.FancyArrow(width/2., height/2., width/5., 0,
length_includes_head=True, width=0,
head_width=height, head_length=height,
overhang=0.2)
return p
Next, a modified version of the add_arrow function from Thomas' answer that uses the FancyArrow patch rather than annotations. This solution might cause weird wrapping like Thomas warned against, but I couldn't figure out how to put the arrow in the legend if the arrow is an annotation.
def add_arrow(line, ax, position=None, direction='right', color=None, label=''):
"""
add an arrow to a line.
line: Line2D object
position: x-position of the arrow. If None, mean of xdata is taken
direction: 'left' or 'right'
color: if None, line color is taken.
label: label for arrow
"""
if color is None:
color = line.get_color()
xdata = line.get_xdata()
ydata = line.get_ydata()
if position is None:
position = xdata.mean()
# find closest index
start_ind = np.argmin(np.absolute(xdata - position))
if direction == 'right':
end_ind = start_ind + 1
else:
end_ind = start_ind - 1
dx = xdata[end_ind] - xdata[start_ind]
dy = ydata[end_ind] - ydata[start_ind]
size = abs(dx) * 5.
x = xdata[start_ind] + (np.sign(dx) * size/2.)
y = ydata[start_ind] + (np.sign(dy) * size/2.)
arrow = patches.FancyArrow(x, y, dx, dy, color=color, width=0,
head_width=size, head_length=size,
label=label,length_includes_head=True,
overhang=0.3, zorder=10)
ax.add_patch(arrow)
Now, a helper function to plot both the arrow and the line. It returns a Line2D object, which is needed for the legend handler we wrote in the first code block
def plot_line_with_arrow(x,y,ax=None,label='',**kw):
if ax is None:
ax = plt.gca()
line = ax.plot(x,y,**kw)[0]
add_arrow(line, ax, label=label)
return line
Finally, we make the plot and update the legend's handler_map with our custom handler.
t = np.linspace(-2, 2, 100)
y = np.sin(t)
line = plot_line_with_arrow(t,y,label='Path', linestyle=':')
plt.gca().set_aspect('equal')
plt.legend(handler_map={patches.FancyArrow :
HandlerLinePatch(patch_func=make_legend_arrow,
linehandle=line)})
plt.show()
Here is the output:
I've found that quiver() works better than arrow() or annotate() when the x and y axes have very different scales. Here's my helper function for plotting a line with arrows:
def plot_with_arrows(ax, x, y, color="g", label="", n_arrows=2):
ax.plot(x, y, rasterized=True, color=color, label=label)
x_range = x.max() - x.min()
y_range = y.max() - y.min()
for i in np.linspace(x.keys().min(), x.keys().max(), n_arrows * 2 + 1).astype(np.int32)[1::2]:
direction = np.array([(x[i+5] - x[i]), (y[i+5] - y[i])])
direction = direction / (np.sqrt(np.sum(np.power(direction, 2)))) * 0.05
direction[0] /= x_range
direction[1] /= y_range
ax.quiver(x[i], y[i], direction[0], direction[1], color=color)
Related
Here is my attempt to change the legend of a barplot from rectangle to square:
import matplotlib.patches as patches
rect1 = patches.Rectangle((0,0),1,1,facecolor='#FF605E')
rect2 = patches.Rectangle((0,0),1,1,facecolor='#64B2DF')
plt.legend((rect1, rect2), ('2016', '2015'))
But when I plot this, I still see rectangles instead of squares:
Any suggestions on how can I do this?
I tried both solutions provided by #ImportanceOfBeingErnest and #furas, here are the results:
#ImportanceOfBeingErnest's solution is the easiest to do:
plt.rcParams['legend.handlelength'] = 1
plt.rcParams['legend.handleheight'] = 1.125
Here is the result:
My final code looks like this:
plt.legend((df.columns[1], df.columns[0]), handlelength=1, handleheight=1) # the df.columns = the legend text
#furas's solution produces this, I don't know why the texts are further away from the rectangles, but I am sure the gap can be changed somehow:
Matplotlib provides the rcParams
legend.handlelength : 2. # the length of the legend lines in fraction of fontsize
legend.handleheight : 0.7 # the height of the legend handle in fraction of fontsize
You can set those within the call to plt.legend()
plt.legend(handlelength=1, handleheight=1)
or using the rcParams at the beginning of your script
import matplotlib
matplotlib.rcParams['legend.handlelength'] = 1
matplotlib.rcParams['legend.handleheight'] = 1
Unfortunately providing equal handlelength=1, handleheight=1 will not give a perfect rectange. It seems handlelength=1, handleheight=1.125 will do the job, but this may depend on the font being used.
An alternative, if you want to use proxy artists may be to use the square markers from the plot/scatter methods.
bar1 = plt.plot([], marker="s", markersize=15, linestyle="", label="2015")
and supply it to the legend, legend(handles=[bar1]). Using this approach needs to have set matplotlib.rcParams['legend.numpoints'] = 1, otherwise two markers would appear in the legend.
Here is a full example of both methods
import matplotlib.pyplot as plt
plt.rcParams['legend.handlelength'] = 1
plt.rcParams['legend.handleheight'] = 1.125
plt.rcParams['legend.numpoints'] = 1
fig, ax = plt.subplots(ncols=2, figsize=(5,2.5))
# Method 1: Set the handlesizes already in the rcParams
ax[0].set_title("Setting handlesize")
ax[0].bar([0,2], [6,3], width=0.7, color="#a30e73", label="2015", align="center")
ax[0].bar([1,3], [3,2], width=0.7, color="#0943a8", label="2016", align="center" )
ax[0].legend()
# Method 2: use proxy markers. (Needs legend.numpoints to be 1)
ax[1].set_title("Proxy markers")
ax[1].bar([0,2], [6,3], width=0.7, color="#a30e73", align="center" )
ax[1].bar([1,3], [3,2], width=0.7, color="#0943a8", align="center" )
b1, =ax[1].plot([], marker="s", markersize=15, linestyle="", color="#a30e73", label="2015")
b2, =ax[1].plot([], marker="s", markersize=15, linestyle="", color="#0943a8", label="2016")
ax[1].legend(handles=[b1, b2])
[a.set_xticks([0,1,2,3]) for a in ax]
plt.show()
producing
It seems they change it long time ago - and now some elements can't be used directly in legend.
Now it needs handler:
Implementing a custom legend handler
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.legend_handler import HandlerPatch
# --- handlers ---
class HandlerRect(HandlerPatch):
def create_artists(self, legend, orig_handle,
xdescent, ydescent, width, height,
fontsize, trans):
x = width//2
y = 0
w = h = 10
# create
p = patches.Rectangle(xy=(x, y), width=w, height=h)
# update with data from oryginal object
self.update_prop(p, orig_handle, legend)
# move xy to legend
p.set_transform(trans)
return [p]
class HandlerCircle(HandlerPatch):
def create_artists(self, legend, orig_handle,
xdescent, ydescent, width, height,
fontsize, trans):
r = 5
x = r + width//2
y = height//2
# create
p = patches.Circle(xy=(x, y), radius=r)
# update with data from oryginal object
self.update_prop(p, orig_handle, legend)
# move xy to legend
p.set_transform(trans)
return [p]
# --- main ---
rect = patches.Rectangle((0,0), 1, 1, facecolor='#FF605E')
circ = patches.Circle((0,0), 1, facecolor='#64B2DF')
plt.legend((rect, circ), ('2016', '2015'),
handler_map={
patches.Rectangle: HandlerRect(),
patches.Circle: HandlerCircle(),
})
plt.show()
Legend reserves place for rectangle and this method doesn't change it so there is so many empty space.
I am trying to draw a series of lines. The lines are all the same length, and randomly switch colors for a random length (blue to orange). I am drawing the lines in blue and then overlaying orange on top. You can see from my picture there are clipped parts of the lines where it is grey. I cannot figure out why this is happening. Also related I believe is that my labels are not moving to a left alignment like they should. Any help is greatly appreciated.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import random
plt.close('all')
fig, ax = plt.subplots(figsize=(15,11))
def label(xy, text):
y = xy[1] - 2
ax.text(xy[0], y, text, ha="left", family='sans-serif', size=14)
def draw_chromosome(start, stop, y, color):
x = np.array([start, stop])
y = np.array([y, y])
line = mlines.Line2D(x , y, lw=10., color=color)
ax.add_line(line)
x = 50
y = 100
chr = 1
for i in range(22):
draw_chromosome(x, 120, y, "#1C2F4D")
j = 0
while j < 120:
print j
length = 1
if random.randint(1, 100) > 90:
length = random.randint(1, 120-j)
draw_chromosome(j, j+length, y, "#FA9B00")
j = j+length+1
label([x, y], "Chromosome%i" % chr)
y -= 3
chr += 1
plt.axis('equal')
plt.axis('off')
plt.tight_layout()
plt.show()
You're only drawing the blue background from x = 50 to x = 120.
Replace this line:
draw_chromosome(x, 120, y, "#1C2F4D")
with this:
draw_chromosome(0, 120, y, "#1C2F4D")
To draw the blue line all the way across.
Alternately, if you also want to move your labels to the left, you can just set x=0 instead of setting it to 50.
I suggest using LineCollection for this. Below is a little helper function I wrote based on the example at http://matplotlib.org/examples/pylab_examples/multicolored_line.html (it looks long, but there is a lot of comments + docstrings)
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import ListedColormap, BoundaryNorm
from matplotlib.ticker import NullLocator
from collections import OrderedDict
def binary_state_lines(ax, chrom_data, xmin=0, xmax=120,
delta_y=3,
off_color = "#1C2F4D",
on_color = "#FA9B00"):
"""
Draw a whole bunch of chromosomes
Parameters
----------
ax : Axes
The axes to draw stuff to
chrom_data : OrderedDict
The chromosome data as a dict, key on the label with a list of pairs
of where the data is 'on'. Data is plotted top-down
xmin, xmax : float, optional
The minimum and maximum limits for the x values
delta_y : float, optional
The spacing between lines
off_color, on_color : color, optional
The colors to use for the the on/off state
Returns
-------
collections : dict
dictionary of the collections added keyed on the label
"""
# base offset
y_val = 0
# make the color map and norm
cmap = ListedColormap([off_color, on_color])
norm = BoundaryNorm([0, 0.5, 1], cmap.N)
# sort out where the text should be
txt_x = (xmax + xmin) / 2
# dictionary to hold the returned artists
ret = dict()
# loop over the input data draw each collection
for label, data in chrom_data.items():
# increment the y offset
y_val += delta_y
# turn the high windows on to alternating
# high/low regions
x = np.asarray(data).ravel()
# assign the high/low state to each one
state = np.mod(1 + np.arange(len(x)), 2)
# deal with boundary conditions to be off
# at start/end
if x[0] > xmin:
x = np.r_[xmin, x]
state = np.r_[0, state]
if x[-1] < xmax:
x = np.r_[x, xmax]
state = np.r_[state, 0]
# make the matching y values
y = np.ones(len(x)) * y_val
# call helper function to create the collection
coll = draw_segments(ax, x, y, state,
cmap, norm)
ret[label] = coll
# set up the axes limits
ax.set_xlim(xmin, xmax)
ax.set_ylim(0, y_val + delta_y)
# turn off x-ticks
ax.xaxis.set_major_locator(NullLocator())
# make the y-ticks be labeled as per the input
ax.yaxis.set_ticks((1 + np.arange(len(chrom_data))) * delta_y)
ax.yaxis.set_ticklabels(list(chrom_data.keys()))
# invert so that the first data is at the top
ax.invert_yaxis()
# turn off the frame and patch
ax.set_frame_on(False)
# return the added artists
return ret
def draw_segments(ax, x, y, state, cmap, norm, lw=10):
"""
helper function to turn boundary edges into the input LineCollection
expects.
Parameters
----------
ax : Axes
The axes to draw to
x, y, state : array
The x edges, the y values and the state of each region
cmap : matplotlib.colors.Colormap
The color map to use
norm : matplotlib.ticker.Norm
The norm to use with the color map
lw : float, optional
The width of the lines
"""
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = LineCollection(segments, cmap=cmap, norm=norm)
lc.set_array(state)
lc.set_linewidth(lw)
ax.add_collection(lc)
return lc
An example:
synthetic_data = OrderedDict()
for j in range(21):
key = 'data {:02d}'.format(j)
synthetic_data[key] = np.cumsum(np.random.randint(1, 10, 20)).reshape(-1, 2)
fig, ax = plt.subplots(tight_layout=True)
binary_state_lines(ax, synthetic_data, xmax=120)
plt.show()
Separating the plotting logic from everything else will make your code easier to maintain and more reusable.
I also took the liberty of moving your labels from between the lines (where they can be ambiguous) to the yaxis tick labels.
I was wondering if it is possible to plot a curve in matplotlib with arrow ticks.
Something like:
from pylab import *
y = linspace(0,10,0.01)
x = cos(y)
plot(x, y, '->')
which should come out with a curve made like this --->---->----> when x increases and like this ---<----<----< whenit decreases (and for y as well, of course).
EDIT:
Furthermore, the arrows should be inclined in the curve's direction (for example, 45 degrees for the y=x function)
It is possible to use the same strategy as in matplotlib streamplot function. Based on the example already given by hitzg:
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.patches as mpatches
def add_arrow_to_line2D(
axes, line, arrow_locs=[0.2, 0.4, 0.6, 0.8],
arrowstyle='-|>', arrowsize=1, transform=None):
"""
Add arrows to a matplotlib.lines.Line2D at selected locations.
Parameters:
-----------
axes:
line: Line2D object as returned by plot command
arrow_locs: list of locations where to insert arrows, % of total length
arrowstyle: style of the arrow
arrowsize: size of the arrow
transform: a matplotlib transform instance, default to data coordinates
Returns:
--------
arrows: list of arrows
"""
if not isinstance(line, mlines.Line2D):
raise ValueError("expected a matplotlib.lines.Line2D object")
x, y = line.get_xdata(), line.get_ydata()
arrow_kw = {
"arrowstyle": arrowstyle,
"mutation_scale": 10 * arrowsize,
}
color = line.get_color()
use_multicolor_lines = isinstance(color, np.ndarray)
if use_multicolor_lines:
raise NotImplementedError("multicolor lines not supported")
else:
arrow_kw['color'] = color
linewidth = line.get_linewidth()
if isinstance(linewidth, np.ndarray):
raise NotImplementedError("multiwidth lines not supported")
else:
arrow_kw['linewidth'] = linewidth
if transform is None:
transform = axes.transData
arrows = []
for loc in arrow_locs:
s = np.cumsum(np.sqrt(np.diff(x) ** 2 + np.diff(y) ** 2))
n = np.searchsorted(s, s[-1] * loc)
arrow_tail = (x[n], y[n])
arrow_head = (np.mean(x[n:n + 2]), np.mean(y[n:n + 2]))
p = mpatches.FancyArrowPatch(
arrow_tail, arrow_head, transform=transform,
**arrow_kw)
axes.add_patch(p)
arrows.append(p)
return arrows
y = np.linspace(0, 100, 200)
x = np.cos(y/5.)
fig, ax = plt.subplots(1, 1)
# print the line and the markers in seperate steps
line, = ax.plot(x, y, 'k-')
add_arrow_to_line2D(ax, line, arrow_locs=np.linspace(0., 1., 200),
arrowstyle='->')
plt.show()
Also refer to this answer.
Try this:
import numpy as np
import matplotlib.pyplot as plt
y = np.linspace(0,100,100)
x = np.cos(y/5.)
# use masked arrays
x1 = np.ma.masked_array(x[:-1], np.diff(x)>=0)
x2 = np.ma.masked_array(x[:-1], np.diff(x)<=0)
# print the line and the markers in seperate steps
plt.plot(x, y, 'k-')
plt.plot(x1, y[:-1], 'k<')
plt.plot(x2, y[:-1], 'k>')
plt.show()
In the figure below, each unit in the x-axis represents a 10mins interval. I would like to customize the labels of x-axis, so that it shows hours, i.e. it displays a ticker every 6 units (60mins). I am new to matplotlib. Could someone help me? Thanks~
Here is the code for the above figure.
x = arange(0, size_x, dx)
y = arange(0, size_y, dy)
X,Y = meshgrid(x, y)
Z = foo(x,y)
pcolor(X, Y, Z, cmap=cm.Reds)
colorbar()
axis([0,size_x-1,0,size_y-1])
show()
There's more than one way to do this.
Let's start out with an example plot:
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
# Generate some data...
x, y = np.mgrid[:141, :101]
z = np.cos(np.hypot(x, y))
# Plot the figure...
plt.pcolormesh(x, y, z, cmap=mpl.cm.Reds)
plt.show()
The simple way to do what you want would be something like this:
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
# Generate some data...
x, y = np.mgrid[:141, :101]
z = np.cos(np.hypot(x, y))
# Plot the figure...
plt.pcolormesh(x, y, z, cmap=mpl.cm.Reds)
# Set the ticks and labels...
ticks = np.arange(x.min(), x.max(), 6)
labels = range(ticks.size)
plt.xticks(ticks, labels)
plt.xlabel('Hours')
plt.show()
The other way involves subclassing matplotlib's locators and tickers.
For your purposes, the example above is fine.
The advantage of making new locators and tickers is that the axis will automatically be scaled into reasonable intervals of the "dx" units you specify. If you're using it as a part of a larger application, it can be worthwhile. For a single plot, it's more trouble than it's worth.
If you really wanted to go that route, though, you'd do something like this:
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
def main():
# Generate some data...
x, y = np.mgrid[:141, :101]
z = np.cos(np.hypot(x, y))
# Plot the figure...
fig, ax = plt.subplots()
ax.pcolormesh(x, y, z, cmap=mpl.cm.Reds)
ax.set_xlabel('Hours')
ax.xaxis.set_major_locator(ScaledLocator(dx=6))
ax.xaxis.set_major_formatter(ScaledFormatter(dx=6))
plt.show()
class ScaledLocator(mpl.ticker.MaxNLocator):
"""
Locates regular intervals along an axis scaled by *dx* and shifted by
*x0*. For example, this would locate minutes on an axis plotted in seconds
if dx=60. This differs from MultipleLocator in that an approriate interval
of dx units will be chosen similar to the default MaxNLocator.
"""
def __init__(self, dx=1.0, x0=0.0):
self.dx = dx
self.x0 = x0
mpl.ticker.MaxNLocator.__init__(self, nbins=9, steps=[1, 2, 5, 10])
def rescale(self, x):
return x / self.dx + self.x0
def inv_rescale(self, x):
return (x - self.x0) * self.dx
def __call__(self):
vmin, vmax = self.axis.get_view_interval()
vmin, vmax = self.rescale(vmin), self.rescale(vmax)
vmin, vmax = mpl.transforms.nonsingular(vmin, vmax, expander = 0.05)
locs = self.bin_boundaries(vmin, vmax)
locs = self.inv_rescale(locs)
prune = self._prune
if prune=='lower':
locs = locs[1:]
elif prune=='upper':
locs = locs[:-1]
elif prune=='both':
locs = locs[1:-1]
return self.raise_if_exceeds(locs)
class ScaledFormatter(mpl.ticker.OldScalarFormatter):
"""Formats tick labels scaled by *dx* and shifted by *x0*."""
def __init__(self, dx=1.0, x0=0.0, **kwargs):
self.dx, self.x0 = dx, x0
def rescale(self, x):
return x / self.dx + self.x0
def __call__(self, x, pos=None):
xmin, xmax = self.axis.get_view_interval()
xmin, xmax = self.rescale(xmin), self.rescale(xmax)
d = abs(xmax - xmin)
x = self.rescale(x)
s = self.pprint_val(x, d)
return s
if __name__ == '__main__':
main()
I'm new to Python (was an IDL user before hand) so I hope that I'm asking this in an understandable way. I've been trying to create a polar plot with x number of bins where the data in the bin is averaged and given a colour associated with that value. This seems to work fine while using the plt.fill command where I can define the bin and then the fill colour. The problem comes when I then try to make a colour bar to go with it. I keep getting errors that state AttributeError: 'Figure' object has no attribute 'autoscale_None'
Any advice would be helpful thanks.
import matplotlib.pyplot as plt
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.pyplot import figure, show, rc, grid
import pylab
r = np.arange(50)/5.
rstep = r[1] - r[0]
theta = np.arange(50)/50.*2.*np.pi
tstep = theta[1] - theta[0]
colorv = np.arange(50)/50.
# force square figure and square axes looks better for polar, IMO
width, height = mpl.rcParams['figure.figsize']
size = min(width, height)
# make a square figure
fig = figure(figsize=(size, size))
ax = fig.add_axes([0.1, 0.1, .8, .8])#, polar=True)
my_cmap = cm.jet
for j in range(len(r)):
rbox = np.array([r[j], r[j], r[j]+ rstep, r[j] + rstep])
for i in range(len(theta)):
thetabox = np.array([theta[i], theta[i] + tstep, theta[i] + tstep, theta[i]])
x = rbox*np.cos(thetabox)
y = rbox*np.sin(thetabox)
plt.fill(x,y, facecolor = my_cmap(colorv[j]))
# Add colorbar, make sure to specify tick locations to match desired ticklabels
cbar = fig.colorbar(fig, ticks=[np.min(colorv), np.max(colorv)])
cb = plt.colorbar()
plt.show()
* here is a slightly better example of my real data, there are holes missing everywhere, so in this example I've just made a big one in a quarter of the circle. When I've tried meshing, the code seems to try to interpolate over these regions.
r = np.arange(50)/50.*7. + 3.
rstep = r[1] - r[0]
theta = np.arange(50)/50.*1.5*np.pi - np.pi
tstep = theta[1] - theta[0]
colorv = np.sin(r/10.*np.pi)
# force square figure and square axes looks better for polar, IMO
width, height = mpl.rcParams['figure.figsize']
size = min(width, height)
# make a square figure
fig = figure(figsize=(size, size))
ax = fig.add_axes([0.1, 0.1, .8, .8])#, polar=True)
my_cmap = cm.jet
for j in range(len(r)):
rbox = np.array([r[j], r[j], r[j]+ rstep, r[j] + rstep])
for i in range(len(theta)):
thetabox = np.array([theta[i], theta[i] + tstep, theta[i] + tstep, theta[i]])
x = rbox*np.cos(thetabox)
y = rbox*np.sin(thetabox)
plt.fill(x,y, facecolor = my_cmap(colorv[j]))
# Add colorbar, make sure to specify tick locations to match desired ticklabels
#cbar = fig.colorbar(fig, ticks=[np.min(colorv), np.max(colorv)])
#cb = plt.colorbar()
plt.show()
And then with a meshing involved...
from matplotlib.mlab import griddata
r = np.arange(50)/5.
rstep = r[1] - r[0]
theta = np.arange(50)/50.*1.5*np.pi - np.pi
tstep = theta[1] - theta[0]
colorv = np.sin(r/10.*np.pi)
# force square figure and square axes looks better for polar, IMO
width, height = mpl.rcParams['figure.figsize']
size = min(width, height)
# make a square figure
fig = figure(figsize=(size, size))
ax = fig.add_axes([0.1, 0.1, .8, .8])#, polar=True)
my_cmap = cm.jet
x = r*np.cos(theta)
y = r*np.sin(theta)
X,Y = np.meshgrid(x,y)
data = griddata(x,y,colorv,X,Y)
cax = plt.contourf(X,Y, data)
plt.colorbar()
# Add colorbar, make sure to specify tick locations to match desired ticklabels
#cbar = fig.colorbar(fig, ticks=[np.min(colorv), np.max(colorv)])
#cb = plt.colorbar()
plt.show()
colorbar needs things to be an instance of ScalarMappable in order to make a colorbar from them.
Because you're manually setting each tile, there's nothing that essentially has a colorbar.
There are a number of ways to fake it from your colormap, but in this case there's a much simpler solution.
pcolormesh does exactly what you want, and will be much faster.
As an example:
import numpy as np
import matplotlib.pyplot as plt
# Linspace makes what you're doing _much_ easier (and includes endpoints)
r = np.linspace(0, 10, 50)
theta = np.linspace(0, 2*np.pi, 50)
fig = plt.figure()
ax = fig.add_subplot(111, projection='polar')
# "Grid" r and theta into 2D arrays (see the docs for meshgrid)
r, theta = np.meshgrid(r, theta)
cax = ax.pcolormesh(theta, r, r, edgecolors='black', antialiased=True)
# We could just call `plt.colorbar`, but I prefer to be more explicit
# and pass in the artist that I want it to extract colors from.
fig.colorbar(cax)
plt.show()
Or, if you'd prefer non-polar axes, as in your example code:
import numpy as np
import matplotlib.pyplot as plt
r = np.linspace(0, 10, 50)
theta = np.linspace(0, 2*np.pi, 50)
# "Grid" r and theta and convert them to cartesian coords...
r, theta = np.meshgrid(r, theta)
x, y = r * np.cos(theta), r * np.sin(theta)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.axis('equal')
cax = ax.pcolormesh(x, y, r, edgecolors='black', antialiased=True)
fig.colorbar(cax)
plt.show()
Note: If you'd prefer the boundary lines a bit less dark, just specify linewidth=0.5 or something similar to pcolormesh.
Finally, if you did want to directly make the colorbar from the colormap in your original code, you'd create an instance of ScalarMappable from it and pass this to colorbar. It's easier than it sounds, but it's a bit verbose.
As an example, in your original code, if you do something like the following:
cax = cm.ScalarMappable(cmap=my_cmap)
cax.set_array(colorv)
fig.colorbar(cax)
It should do what you want.
So I've found a workaround. Since I know of a region where I definitely won't have data, I've plotted some there. I've made sure that the data covers the entire range of what I'm potting. I then cover it up (this region was going to be covered anyway, it shows where the "earth" is located). Now I can go ahead and use plt.fill as I had originally and use the colour bar from the randomly potted data. I know this isn't probably the correct way, but it works and doesn't try to interpolate my data.
Thanks so much for helping get this sorted. and if you know of a better way, I'd be happy to hear it!
hid = plt.pcolormesh(X,Y, data, antialiased=True)
#here we cover up the region that we just plotted in
r3 = [1 for i in range(360)]
theta3 = np.arange(360)*np.pi/180.
plt.fill(theta3, r3, 'w')
#now we can go through and fill in all the regions
for j in range(len(r)):
rbox = np.array([r[j], r[j], r[j]+ rstep, r[j] + rstep])
for i in range(len(theta)):
thetabox = np.array([theta[i], theta[i] + tstep, theta[i] + tstep, theta[i]])
x = rbox*np.cos(thetabox)
y = rbox*np.sin(thetabox)
colorv = np.sin(r[j]/10.*np.pi)
plt.fill(thetabox,rbox, facecolor = my_cmap(colorv))
#And now we can plot the color bar that fits the data Tada :)
plt.colorbar()
plt.show()