I modified a function based on a nearest neighbor algorithm that takes the parameters x and y from a list of randomly generated values to plot points on a graph based, but the edges are being plotted out wrong and I can't figure out why.
Here is the original algorithm I based my code on
import matplotlib.pyplot as plt
import numpy as np
xy = np.random.rand(40, 2)
plt.figure(figsize=(8,8))
dist_sq = np.sum((xy[:,np.newaxis,:] - xy[np.newaxis,:,:]) ** 2, axis=-1)
#nearest = np.argsort(dist_sq, axis=1)
K = 1
nearest_partition = np.argpartition(dist_sq, K+1, axis=1)
plt.figure(figsize=(8,8))
#plt.scatter(xy[:,0], xy[:,1], s=40)
for i in range(xy.shape[0]):
for j in nearest_partition[i, :K+1]:
plt.plot(*zip(xy[j], xy[i]), color='black')
and here is my function
import matplotlib.pyplot as plt
import numpy as np
def nearest_neighbor_graph(x, y, k, pcolor='blue', psize=20, ecolor='black', figsize=(6,6)):
plt.figure(figsize=(6, 6))
plt.scatter(x, y, psize, color = pcolor)
dist_sq = np.sum((xy[:,np.newaxis,:] - xy[np.newaxis,:,:]) ** 2, axis=-1)
nearest = np.argsort(dist_sq, axis=1)
#nearest_partition = np.argpartition(dist_sq, k + 1, axis=1)
for i in range(xy.shape[0]):
for j in nearest[i, :k + 1]:
plt.plot(*zip(xy[i], xy[j]), color= ecolor)
np.random.seed(10)
xy = np.random.rand(2, 40)
nearest_neighbor_graph(xy[0], xy[1], 1, ecolor='red', psize=50)
nearest_neighbor_graph(xy[0], xy[1], 2, pcolor='green', psize=100, ecolor='black')
In my function, when I get to the nearest_partition line I get an error so I use nearest instead in its place in the bottom loop
for i in range(xy.shape[0]):
for j in nearest[i, :k + 1]:
plt.plot(*zip(xy[i], xy[j]), color= ecolor)
Any help would be greatly appreciated, thanks
You were really close. The error message was the clue:
ValueError: kth(=2) out of bounds (2)
checked the array size with np.shape() and it was 40x40, so then checked the syntax which was incorrect. You can see the corrected line.
so this works:
import matplotlib.pyplot as plt
import numpy as np
def nearest_neighbor_graph(x, y, k, pcolor='blue', psize=20, ecolor='black', figsize=(6,6)):
plt.figure(figsize=(6, 6))
plt.scatter(x, y, psize, color = pcolor)
dist_sq = np.sum((xy[:,np.newaxis,:] - xy[np.newaxis,:,:]) ** 2, axis=-1)
# nearest = np.argsort(dist_sq, axis=1)
nearest_partition = np.argpartition(dist_sq, kth= 1, axis=1)
for i in range(xy.shape[0]):
for j in nearest_partition[i, :k + 1]:
plt.plot(*zip(xy[i], xy[j]), color= ecolor)
np.random.seed(10)
xy = np.random.rand(2, 40)
nearest_neighbor_graph(xy[0], xy[1], 1, ecolor='red', psize=50)
nearest_neighbor_graph(xy[0], xy[1], 2, pcolor='green', psize=100, ecolor='black')
Here are the resulting graphs:
Related
How I can fill the common area under both the curve?
import matplotlib.pyplot as plt
import numpy as np
import scipy.special as sp
x = np.linspace(-4, 4, num=1000)
r = abs(x)
zeta = 1.0
psi_STO = (zeta**3 / np.pi)**(0.5) * np.exp(-zeta * r)
plt.figure(figsize=(4, 3))
plt.plot(x, psi_STO, color="C0")
plt.plot(x + 3, psi_STO, color="C0")
plt.show()
If I use:
plt.fill_betweenx(psi_STO, -1, 4, color="C1")
I am getting a plot as:
You can use fill_between. As your Xs are not aligned, you need to make a bit of calculations first to find the common range. This will depend on the number of points in the linspace. Here, I computed it manually: as the shift is of 3, there are 375 points difference (250 per unit).
import matplotlib.pyplot as plt
import numpy as np
import scipy.special as sp
x = np.linspace(-4, 4, num=1000)
r = abs(x)
zeta = 1.0
psi_STO = (zeta**3 / np.pi)**(0.5) * np.exp(-zeta * r)
plt.figure(figsize=(4, 3))
plt.plot(x, psi_STO, color="C0")
plt.plot(x + 3, psi_STO, color="C0")
x_common = (x+3)[:-375]
min_common = np.min([psi_STO[:-375], psi_STO[375:]], axis=0)
plt.plot(x_common, min_common, color='r')
plt.fill_between(x_common,min_common, color="#FFB0B0")
plt.show()
output:
This?
dx = 3 # x shift
di = int(dx/(x[1]-x[0])) # index shift
plt.fill_between(x[di:], np.minimum(psi_STO[:-di], psi_STO[di:]))
I have the following code:
import numpy as np
import matplotlib.pyplot as plt
x = np.linspace(-np.pi/2, np.pi/2, 30)
y = np.linspace(-np.pi/2, np.pi/2, 30)
x,y = np.meshgrid(x,y)
z = np.sin(x**2+y**2)[:-1,:-1]
fig,ax = plt.subplots()
ax.pcolormesh(x,y,z)
Which gives this image:
Now lets say I want to highlight the edge certain grid boxes:
highlight = (z > 0.9)
I could use the contour function, but this would result in a "smoothed" contour. I just want to highlight the edge of a region, following the edge of the grid boxes.
The closest I've come is adding something like this:
highlight = np.ma.masked_less(highlight, 1)
ax.pcolormesh(x, y, highlight, facecolor = 'None', edgecolors = 'w')
Which gives this plot:
Which is close, but what I really want is for only the outer and inner edges of that "donut" to be highlighted.
So essentially I am looking for some hybrid of the contour and pcolormesh functions - something that follows the contour of some value, but follows grid bins in "steps" rather than connecting point-to-point. Does that make sense?
Side note: In the pcolormesh arguments, I have edgecolors = 'w', but the edges still come out to be blue. Whats going on there?
EDIT:
JohanC's initial answer using add_iso_line() works for the question as posed. However, the actual data I'm using is a very irregular x,y grid, which cannot be converted to 1D (as is required for add_iso_line().
I am using data which has been converted from polar coordinates (rho, phi) to cartesian (x,y). The 2D solution posed by JohanC does not appear to work for the following case:
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage
def pol2cart(rho, phi):
x = rho * np.cos(phi)
y = rho * np.sin(phi)
return(x, y)
phi = np.linspace(0,2*np.pi,30)
rho = np.linspace(0,2,30)
pp, rr = np.meshgrid(phi,rho)
xx,yy = pol2cart(rr, pp)
z = np.sin(xx**2 + yy**2)
scale = 5
zz = ndimage.zoom(z, scale, order=0)
fig,ax = plt.subplots()
ax.pcolormesh(xx,yy,z[:-1, :-1])
xlim = ax.get_xlim()
ylim = ax.get_ylim()
xmin, xmax = xx.min(), xx.max()
ymin, ymax = yy.min(), yy.max()
ax.contour(np.linspace(xmin,xmax, zz.shape[1]) + (xmax-xmin)/z.shape[1]/2,
np.linspace(ymin,ymax, zz.shape[0]) + (ymax-ymin)/z.shape[0]/2,
np.where(zz < 0.9, 0, 1), levels=[0.5], colors='red')
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)
This post shows a way to draw such lines. As it is not straightforward to adapt to the current pcolormesh, the following code demonstrates a possible adaption.
Note that the 2d versions of x and y have been renamed, as the 1d versions are needed for the line segments.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
x = np.linspace(-np.pi / 2, np.pi / 2, 30)
y = np.linspace(-np.pi / 2, np.pi / 2, 30)
xx, yy = np.meshgrid(x, y)
z = np.sin(xx ** 2 + yy ** 2)[:-1, :-1]
fig, ax = plt.subplots()
ax.pcolormesh(x, y, z)
def add_iso_line(ax, value, color):
v = np.diff(z > value, axis=1)
h = np.diff(z > value, axis=0)
l = np.argwhere(v.T)
vlines = np.array(list(zip(np.stack((x[l[:, 0] + 1], y[l[:, 1]])).T,
np.stack((x[l[:, 0] + 1], y[l[:, 1] + 1])).T)))
l = np.argwhere(h.T)
hlines = np.array(list(zip(np.stack((x[l[:, 0]], y[l[:, 1] + 1])).T,
np.stack((x[l[:, 0] + 1], y[l[:, 1] + 1])).T)))
lines = np.vstack((vlines, hlines))
ax.add_collection(LineCollection(lines, lw=1, colors=color))
add_iso_line(ax, 0.9, 'r')
plt.show()
Here is an adaption of the second answer, which can work with only 2d arrays:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from scipy import ndimage
x = np.linspace(-np.pi / 2, np.pi / 2, 30)
y = np.linspace(-np.pi / 2, np.pi / 2, 30)
x, y = np.meshgrid(x, y)
z = np.sin(x ** 2 + y ** 2)
scale = 5
zz = ndimage.zoom(z, scale, order=0)
fig, ax = plt.subplots()
ax.pcolormesh(x, y, z[:-1, :-1] )
xlim = ax.get_xlim()
ylim = ax.get_ylim()
xmin, xmax = x.min(), x.max()
ymin, ymax = y.min(), y.max()
ax.contour(np.linspace(xmin,xmax, zz.shape[1]) + (xmax-xmin)/z.shape[1]/2,
np.linspace(ymin,ymax, zz.shape[0]) + (ymax-ymin)/z.shape[0]/2,
np.where(zz < 0.9, 0, 1), levels=[0.5], colors='red')
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)
plt.show()
I'll try to refactor add_iso_line method in order to make it more clear an open for optimisations. So, at first, there comes a must-do part:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
x = np.linspace(-np.pi/2, np.pi/2, 30)
y = np.linspace(-np.pi/2, np.pi/2, 30)
x, y = np.meshgrid(x,y)
z = np.sin(x**2+y**2)[:-1,:-1]
fig, ax = plt.subplots()
ax.pcolormesh(x,y,z)
xlim, ylim = ax.get_xlim(), ax.get_ylim()
highlight = (z > 0.9)
Now highlight is a binary array that looks like this:
After that we can extract indexes of True cells, look for False neighbourhoods and identify positions of 'red' lines. I'm not comfortable enough with doing it in a vectorised manner (like here in add_iso_line method) so just using simple loop:
lines = []
cells = zip(*np.where(highlight))
for x, y in cells:
if x == 0 or highlight[x - 1, y] == 0: lines.append(([x, y], [x, y + 1]))
if x == highlight.shape[0] or highlight[x + 1, y] == 0: lines.append(([x + 1, y], [x + 1, y + 1]))
if y == 0 or highlight[x, y - 1] == 0: lines.append(([x, y], [x + 1, y]))
if y == highlight.shape[1] or highlight[x, y + 1] == 0: lines.append(([x, y + 1], [x + 1, y + 1]))
And, finally, I resize and center coordinates of lines in order to fit with pcolormesh:
lines = (np.array(lines) / highlight.shape - [0.5, 0.5]) * [xlim[1] - xlim[0], ylim[1] - ylim[0]]
ax.add_collection(LineCollection(lines, colors='r'))
plt.show()
In conclusion, this is very similar to JohanC solution and, in general, slower. Fortunately, we can reduce amount of cells significantly, extracting contours only using python-opencv package:
import cv2
highlight = highlight.astype(np.uint8)
contours, hierarchy = cv2.findContours(highlight, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
cells = np.vstack(contours).squeeze()
This is an illustration of cells being checked:
I have a 3D polygon plot and want to smooth the plot on the y axis (i.e. I want it to look like 'slices of a surface plot').
Consider this MWE (taken from here):
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.collections import PolyCollection
import matplotlib.pyplot as plt
from matplotlib import colors as mcolors
import numpy as np
from scipy.stats import norm
fig = plt.figure()
ax = fig.gca(projection='3d')
xs = np.arange(-10, 10, 2)
verts = []
zs = [0.0, 1.0, 2.0, 3.0]
for z in zs:
ys = np.random.rand(len(xs))
ys[0], ys[-1] = 0, 0
verts.append(list(zip(xs, ys)))
poly = PolyCollection(verts, facecolors=[mcolors.to_rgba('r', alpha=0.6),
mcolors.to_rgba('g', alpha=0.6),
mcolors.to_rgba('b', alpha=0.6),
mcolors.to_rgba('y', alpha=0.6)])
poly.set_alpha(0.7)
ax.add_collection3d(poly, zs=zs, zdir='y')
ax.set_xlabel('X')
ax.set_xlim3d(-10, 10)
ax.set_ylabel('Y')
ax.set_ylim3d(-1, 4)
ax.set_zlabel('Z')
ax.set_zlim3d(0, 1)
plt.show()
Now, I want to replace the four plots with normal distributions (to ideally form continuous lines).
I have created the distributions here:
def get_xs(lwr_bound = -4, upr_bound = 4, n = 80):
""" generates the x space betwee lwr_bound and upr_bound so that it has n intermediary steps """
xs = np.arange(lwr_bound, upr_bound, (upr_bound - lwr_bound) / n) # x space -- number of points on l/r dimension
return(xs)
xs = get_xs()
dists = [1, 2, 3, 4]
def get_distribution_params(list_):
""" generates the distribution parameters (mu and sigma) for len(list_) distributions"""
mus = []
sigmas = []
for i in range(len(dists)):
mus.append(round((i + 1) + 0.1 * np.random.randint(0,10), 3))
sigmas.append(round((i + 1) * .01 * np.random.randint(0,10), 3))
return mus, sigmas
mus, sigmas = get_distribution_params(dists)
def get_distributions(list_, xs, mus, sigmas):
""" generates len(list_) normal distributions, with different mu and sigma values """
distributions = [] # distributions
for i in range(len(list_)):
x_ = xs
z_ = norm.pdf(xs, loc = mus[i], scale = sigmas[0])
distributions.append(list(zip(x_, z_)))
#print(x_[60], z_[60])
return distributions
distributions = get_distributions(list_ = dists, xs = xs, mus = mus, sigmas = sigmas)
But adding them to the code (with poly = PolyCollection(distributions, ...) and ax.add_collection3d(poly, zs=distributions, zdir='z') throws a ValueError (ValueError: input operand has more dimensions than allowed by the axis remapping) I cannot resolve.
The error is caused by passing distributions to zs where zs expects that when verts in PolyCollection has shape MxNx2 the object passed to zs has shape M. So when it reaches this check
cpdef ndarray broadcast_to(ndarray array, shape):
# ...
if array.ndim < len(shape):
raise ValueError(
'input operand has more dimensions than allowed by the axis '
'remapping')
# ...
in the underlying numpy code, it fails. I believe this occurs because the number of dimensions expected (array.ndim) is less than the number of dimensions of zs (len(shape)). It is expecting an array of shape (4,) but receives an array of shape (4, 80, 2).
This error can be resolved by using an array of the correct shape - e.g. zs from the original example or dists from your code. Using zs=dists and adjusting the axis limits to [0,5] for x, y, and z gives
This looks a bit odd for two reasons:
There is a typo in z_ = norm.pdf(xs, loc = mus[i], scale = sigmas[0]) which gives all the distributions the same sigma, it should be z_ = norm.pdf(xs, loc = mus[i], scale = sigmas[i])
The viewing geometry: the distributions have the positive xz plane as their base, this is also the plane we are looking through.
Changing the viewing geometry via ax.view_init will yield a clearer plot:
Edit
Here is the complete code which generates the plot shown,
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.collections import PolyCollection
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
from scipy.stats import norm
np.random.seed(8)
def get_xs(lwr_bound = -4, upr_bound = 4, n = 80):
return np.arange(lwr_bound, upr_bound, (upr_bound - lwr_bound) / n)
def get_distribution_params(list_):
mus = [round((i+1) + 0.1 * np.random.randint(0,10), 3) for i in range(len(dists))]
sigmas = [round((i+1) * .01 * np.random.randint(0,10), 3) for i in range(len(dists))]
return mus, sigmas
def get_distributions(list_, xs, mus, sigmas):
return [list(zip(xs, norm.pdf(xs, loc=mus[i], scale=sigmas[i] if sigmas[i] != 0.0
else 0.1))) for i in range(len(list_))]
dists = [1, 2, 3, 4]
xs = get_xs()
mus, sigmas = get_distribution_params(dists)
distributions = get_distributions(dists, xs, mus, sigmas)
fc = [mcolors.to_rgba('r', alpha=0.6), mcolors.to_rgba('g', alpha=0.6),
mcolors.to_rgba('b', alpha=0.6), mcolors.to_rgba('y', alpha=0.6)]
poly = PolyCollection(distributions, fc=fc)
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.add_collection3d(poly, zs=np.array(dists).astype(float), zdir='z')
ax.view_init(azim=115)
ax.set_zlim([0, 5])
ax.set_ylim([0, 5])
ax.set_xlim([0, 5])
I based it off the code you provide in the question, but made some modifications for brevity and to be more consistent with the usual styling.
Note - The example code you have given will fail depending on the np.random.seed(), in order to ensure it works I have added a check in the call to norm.pdf which ensures the scale is non-zero: scale = sigma[i] if sigma[i] != 0.0 else 0.1.
Using ax.add_collection3d(poly, zs=dists, zdir='z') instead of ax.add_collection3d(poly, zs=distributions, zdir='z') should fix the issue.
Additionally, you might want to replace
def get_xs(lwr_bound = -4, upr_bound = 4, n = 80):
""" generates the x space betwee lwr_bound and upr_bound so that it has n intermediary steps """
xs = np.arange(lwr_bound, upr_bound, (upr_bound - lwr_bound) / n) # x space -- number of points on l/r dimension
return(xs)
xs = get_xs()
by
xs = np.linspace(-4, 4, 80)
Also, I believe the scale = sigmas[0] should actually be scale = sigmas[i] in the line
z_ = norm.pdf(xs, loc = mus[i], scale = sigmas[0])
Finally, I believe you should adjust the xlim, ylim and zlim appropriatly, as you swapped the y and z dimensions of the plot and changed its scales when comparing to the reference code.
I want to start the curve with one color and progressively blend into another color until the end. The following function in my MCVE works, but surely, there has to be a better way I haven't found out about, yet?!
import numpy as np
import matplotlib.pyplot as plt
def colorlist(color1, color2, num):
"""Generate list of num colors blending from color1 to color2"""
result = [np.array(color1), np.array(color2)]
while len(result) < num:
temp = [result[0]]
for i in range(len(result)-1):
temp.append(np.sqrt((result[i]**2+result[i+1]**2)/2))
temp.append(result[i+1])
result = temp
indices = np.linspace(0, len(result)-1, num).round().astype(int)
return [result[i] for i in indices]
x = np.linspace(0, 2*np.pi, 100)
y = np.sin(x)
colors = colorlist((1, 0, 0), (0, 0, 1), len(x))
for i in range(len(x)-1):
xi = x[i:i+1+1]
yi = y[i:i+1+1]
ci = colors[i]
plt.plot(xi, yi, color=ci, linestyle='solid', linewidth='10')
plt.show()
Not sure what "better way" refers to. A solution with less code, which would draw faster is the use of a LineCollection together with a colormap.
A colormap can be defined by two colors and any colors in between are automatically interpolated.
cmap = matplotlib.colors.LinearSegmentedColormap.from_list("", [(1, 0, 0), (0, 0, 1)])
A LineCollection can be used to plot a lot of lines at once. Being a ScalarMappable it can use a colormap to colorize each line differently according to some array - in this case one may just use the x values for that purpose.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import LinearSegmentedColormap
x = np.linspace(0, 2*np.pi, 100)
y = np.sin(x)
cmap = LinearSegmentedColormap.from_list("", [(1, 0, 0), (0, 0, 1)])
points = np.array([x, y]).T.reshape(-1,1,2)
segments = np.concatenate([points[:-1],points[1:]], axis=1)
lc = LineCollection(segments, cmap=cmap, linewidth=10)
lc.set_array(x)
plt.gca().add_collection(lc)
plt.gca().autoscale()
plt.show()
The drawback of this solution as can be see in the picture is that the individual lines are not well connected.
So to circumvent this, one may plot those points overlapping, using
segments = np.concatenate([points[:-2],points[1:-1], points[2:]], axis=1)
In the above the color is linearly interpolated between the two given colors. The plot therefore looks different than the one from the question using some custom interpolation.
To obtain the same colors as in the question, you may use the same function to create the colors used in the colormap for the LineCollection. If the aim is to simplify this function you may directly calculate the values as the square root of the color difference in the channels.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import LinearSegmentedColormap
x = np.linspace(0, 2*np.pi, 100)
y = np.sin(x)
def colorlist2(c1, c2, num):
l = np.linspace(0,1,num)
a = np.abs(np.array(c1)-np.array(c2))
m = np.min([c1,c2], axis=0)
s = np.sign(np.array(c2)-np.array(c1)).astype(int)
s[s==0] =1
r = np.sqrt(np.c_[(l*a[0]+m[0])[::s[0]],(l*a[1]+m[1])[::s[1]],(l*a[2]+m[2])[::s[2]]])
return r
cmap = LinearSegmentedColormap.from_list("", colorlist2((1, 0, 0), (0, 0, 1),100))
points = np.array([x, y]).T.reshape(-1,1,2)
segments = np.concatenate([points[:-2],points[1:-1], points[2:]], axis=1)
lc = LineCollection(segments, cmap=cmap, linewidth=10)
lc.set_array(x)
plt.gca().add_collection(lc)
plt.gca().autoscale()
plt.show()
In response to a comment above: If you want to change the color depending on the y value, you can use the following code:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import LinearSegmentedColormap
x = np.linspace(0, 2 * np.pi, 100)
y = np.sin(x)
ynorm = (y - y.min()) / (y.max() - y.min())
def colorlist2(c1, c2, num):
l = np.linspace(0, 1, num)
a = np.abs(np.array(c1) - np.array(c2))
m = np.min([c1, c2], axis=0)
s = np.sign(np.array(c2) - np.array(c1)).astype(int)
s[s == 0] = 1
r = np.sqrt(np.c_[(l * a[0] + m[0])[::s[0]],
(l * a[1] + m[1])[::s[1]], (l * a[2] + m[2])[::s[2]]])
return r
cmap = LinearSegmentedColormap.from_list(
"", colorlist2((1, 0, 0), (0, 0, 1), 100))
colors = [cmap(k) for k in ynorm[:-1]]
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-2], points[1:-1], points[2:]], axis=1)
lc = LineCollection(segments, colors=colors, linewidth=10)
lc.set_array(x)
plt.gca().add_collection(lc)
plt.gca().autoscale()
plt.show()
This will output this graph:
Graph with color depending on y value
This code:
def complex_to_rgb(complex_data, invert=False):
from numpy import angle, max, pi, sin, zeros
phase = angle(complex_data)
amplitude = abs(complex_data)
amplitude = amplitude/max(max(amplitude))
A = zeros((complex_data.shape[0], complex_data.shape[1], 3))
A[:,:,0] = .5*(sin(phase)+1)*amplitude
A[:,:,1] = .5*(sin(phase+pi/2)+1)*amplitude
A[:,:,2] = .5*(-sin(phase)+1)*amplitude
if(invert):
return 1-A
else:
return A
import numpy as np
from matplotlib.pyplot import figure
N = 1024
x = np.linspace(-1, 1, N)
y = np.linspace(-1, 1, N)
X,Y = np.meshgrid(x,y)
R = np.sqrt(X*X + Y*Y)
PHI = np.arctan2(Y, X)
fig = figure()
ax = fig.add_subplot(212, polar=True)
ax.imshow(complex_to_rgb(R*np.exp(1j*PHI) * (R<1), invert=True))
ax.set_xticks([-.5, 0, np.pi/2, np.pi, 3*np.pi/2])
ax.set_yticks([0, N/3, 2*N/3, N])
ax.set_xticklabels(['', '$0$', r'$\pi/2$', r'$\pi$', r'$3\pi/2$'])
ax.set_yticklabels([])
fig.show()
Generates a nice HSV legend plot. Now I'd like to remove the -.5 xtick, but that seems to mess everything up. Anyone know how to fix this? I already reported it as a bug
As described in the bug report, I can place the radial axis anywhere I want by specifying an explicit extent to imshow. Additionally, rgrids can be used to fix the angle of the tick labels.
def complex_to_rgb(complex_data, invert=False):
from numpy import angle, max, pi, sin, zeros
phase = angle(complex_data)
amplitude = abs(complex_data)
amplitude = amplitude/max(max(amplitude))
A = zeros((complex_data.shape[0], complex_data.shape[1], 3))
A[:,:,0] = .5*(sin(phase)+1)*amplitude
A[:,:,1] = .5*(sin(phase+pi/2)+1)*amplitude
A[:,:,2] = .5*(-sin(phase)+1)*amplitude
if(invert):
return 1-A
else:
return A
import numpy as np
from matplotlib.pyplot import figure
N = 1024
x = np.linspace(-1, 1, N)
y = np.linspace(-1, 1, N)
X,Y = np.meshgrid(x,y)
R = np.sqrt(X*X + Y*Y)
PHI = np.arctan2(Y, X)
fig = figure()
ax = fig.add_subplot(111, polar=True)
ax.imshow(complex_to_rgb(R*np.exp(1j*PHI) * (R<1), invert=True), extent=[0,2*np.pi, 0,1024])
ax.set_rgrids([1,N/3,2*N/3], angle=45)
ax.set_xticks([0, np.pi/2, np.pi, 3*np.pi/2])
ax.set_yticks([0, N/3, 2*N/3, N])
ax.set_xticklabels([r'$0$', r'$\pi/2$', r'$\pi$', r'$3\pi/2$'])
ax.set_yticklabels([r'0', r'$1/3$', r'$2/3$', '1'])
fig.show()
Which results in: