Related
Trying to plot two separate animations, i.e. in different windows as separate figures. Running this code for me rightly creates two windows, but animates the data on the second figure at the same time. Closing figure 1 results in only the intended data for figure 2 being animated, removing the overlap from the data intended for figure 1. Closing figure 2 results in only the intended data for figure 1 being animated, removing the overlap from the data intended for figure 2.
Minimum code below:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
dx, dv, N, Nb, decp = 2, 1.5, 100, 12, int(1)
Pd = np.zeros([N + 1, 2 * Nb])
Vd = np.zeros([N + 1, 2 * Nb])
Pd[:, 1] = 4
Vd[:, 3] = 2
t = np.zeros(N + 1)
t[0] = 0
for i in range(0, N):
t[i + 1] = (i + 1) * 0.1
Px = []
for i in range(0, (2 * Nb)):
PX = dx * (-Nb + i) / 4
Px.append(PX)
lblx = []
for i in range(0, int((Nb / 2) + 1)):
if i == (Nb / 4):
LBL = r"$\mu_x$"
lblx.append(LBL)
else:
LBL = r"${0}\sigma_x$".format(-(Nb / 4) + i)
lblx.append(LBL)
Pv = []
for i in range(0, (2 * Nb)):
PV = dv * (-Nb + i) / 4
Pv.append(PV)
lblv = []
for i in range(0, int((Nb / 2) + 1)):
if i == (Nb / 4):
LBL = r"$\mu_v$"
lblv.append(LBL)
else:
LBL = r"${0}\sigma_v$".format(-(Nb / 4) + i)
lblv.append(LBL)
fig1 = plt.figure(figsize=(8,6))
def animatex(i):
fig1.clear()
plt.bar(Px, Pd[i, :], width = dx / 4, align = 'edge', color = 'b', \
label = 't = {} seconds'.format(round(t[i], decp)))
s_ticks = np.arange(-3 * dx, (3 + 1) * dx, dx)
plt.xticks(s_ticks, lblx)
plt.ylim(0, np.max(Pd))
plt.xlim(-3 * dx, 3 * dx)
plt.legend()
plt.draw()
anix = FuncAnimation(fig1, animatex, repeat = True, interval = 200, frames = N + 1)
fig2 = plt.figure(figsize=(8,6))
def animatev(i):
fig2.clear()
plt.bar(Pv, Vd[i, :], width = dv / 4, align = 'edge', color = 'b', \
label = 't = {} seconds'.format(round(t[i], decp)))
s_ticks = np.arange(-3 * dv, (3 + 1) * dv, dv)
plt.xticks(s_ticks, lblv)
plt.ylim(0, np.max(Vd))
plt.xlim(-3 * dv, 3 * dv)
plt.legend()
plt.draw()
aniv = FuncAnimation(fig2, animatev, repeat = True, interval = 200, frames = N + 1)
plt.show()
As is probably clear, they are two bar plots, with different vertical and horizontal dimensions. I've seen some solutions for these kinds of problems where the data shares an axis through a shared variable, but here they are not (as can be seen).
For this minimum code, the solution involves having the two bars, one in Pd and the other in Vd, being on their respective intended figures, not both on the second figure.
Let me know if there are any issues with the information here i.e. minimal code requirements not met, more information etc. and I will update.
Ignore any wayward writing style, it is not relevant.
Simplifying your code:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
dx, dv, N, Nb, decp = 2, 1.5, 10, 12, int(1)
Px = np.arange(Nb)
Pd = np.random.randn(N, Nb)
Vd = np.random.randn(N, Nb)
fig1, ax1 = plt.subplots(figsize=(8, 6))
def animatex(i):
ax1.clear()
ax1.bar(Px, Pd[i, :], width=dx / 4, align='edge', color='b')
anix = FuncAnimation(fig1, animatex, repeat=True, interval=200, frames=N)
fig2, ax2 = plt.subplots(figsize=(8, 6))
def animatev(i):
ax2.clear()
ax2.bar(Px, Vd[i, :], width = dv / 4, align='edge', color='b')
aniv = FuncAnimation(fig2, animatev, repeat=True, interval=200, frames=N)
plt.show()
works fine for me. You can add the esthetic/data details back in...
I'm looking to have a main image upon which I draw either spirals, ellipses etc with variables that change the shape on the imposed drawing. The main image also needs to have a contrast variable.
My code currently looks like this;
###############################################BASIC FIGURE PLOT####################################
plt.figure(figsize=(24,24))
#interact
def spiral(Spiral=False,n=2000,x1=50,y1=50,z1=50,k1=300):
if Spiral == False:
x = 0;
y = 0;
plt.scatter(x,y,s = 3, c = 'black');
else:
angle = np.linspace(x1,y1*1*np.pi, n)
radius = np.linspace(z1,k1,n)
x = radius * np.cos(angle) + 150
y = radius * np.sin(angle) + 150
plt.scatter(x,y,s = 3, c = 'black');
#interact
def contrast(vuc=(0.2,1,0.01),vlc=(0.1,1,0.01)):
vu = np.quantile(qphi, vuc);
vl = np.quantile(qphi, vlc);
print("upper =",vu, " lower=",vl);
plt.imshow(qphi, origin='lower',vmin=vl,vmax=vu);
plt.show()
This produces two plots;
visible here
One plot which creates a spiral I can edit freely and one plot that is the main image with variable contrast.
Any advise on how to combine the two plots would be much appreciated; Thank you!
There are several ways to approach controlling a matplotlib plot using ipywidgets. Below I've created the output I think you're looking for using each of the options. The methods are listed in what feels like the natural order of discovery, however, I would recommend trying them in this order: 4, 2, 1, 3
Approach 1 - inline backend
If you use %matplotlib inline then matplotlib figures will not be interactive and you will need to recreate the entire plot every time
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from ipywidgets import interact
# load fake image data
from matplotlib import cbook
img = plt.imread(cbook.get_sample_data("grace_hopper.jpg")).mean(axis=-1)
#interact
def graph(
Spiral=True,
n=2000,
x1=50,
y1=50,
z1=50,
k1=300,
vlc=(0.1, 1, 0.01),
vuc=(0.1, 1, 0.01),
):
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12,5))
if Spiral == False:
x = 0
y = 0
else:
angle = np.linspace(x1, y1 * 1 * np.pi, n)
radius = np.linspace(z1, k1, n)
x = radius * np.cos(angle) + 150
y = radius * np.sin(angle) + 150
ax1.scatter(x, y, s=3, color="k")
vu = np.quantile(img, vuc)
vl = np.quantile(img, vlc)
ax2.imshow(img, vmin=vl, vmax=vu)
Approach 2 - interactive backend + cla
You can use one of the interactive maptlotlib backends to avoid having to completely regenerate the figure every time you change. To do this the first approach is to simply clear the axes everytime the sliders change using the cla method.
This will work with either %matplotlib notebook or %matplotlib ipympl. The former will only work in jupyter notebook and the latter will work in both jupyter notebook and juptyerlab. (Installation info for ipympl here: https://github.com/matplotlib/ipympl#installation)
%matplotlib ipympl
import matplotlib.pyplot as plt
import numpy as np
from ipywidgets import interact, interactive, interactive_output
# load fake image data
from matplotlib import cbook
img = plt.imread(cbook.get_sample_data("grace_hopper.jpg")).mean(axis=-1)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12,5))
#interact
def graph(
Spiral=True,
n=2000,
x1=50,
y1=50,
z1=50,
k1=300,
vlc=(0.1, 1, 0.01),
vuc=(0.1, 1, 0.01),
):
ax1.cla()
ax2.cla()
if Spiral == False:
x = 0
y = 0
else:
angle = np.linspace(x1, y1 * 1 * np.pi, n)
radius = np.linspace(z1, k1, n)
x = radius * np.cos(angle) + 150
y = radius * np.sin(angle) + 150
ax1.scatter(x, y, s=3, color="k")
vu = np.quantile(img, vuc)
vl = np.quantile(img, vlc)
ax2.imshow(img, vmin=vl, vmax=vu)
Approach 3 - interactive backend + set_data
Totally clearing the axes can be inefficient when you are plotting larger datasets or have some parts of the plot that you want to persist from one interaction to the next. So you can instead use the set_data and set_offsets methods to update what you have already drawn.
%matplotlib ipympl
import matplotlib.pyplot as plt
import numpy as np
from ipywidgets import interact, interactive, interactive_output
# load fake image data
from matplotlib import cbook
img = plt.imread(cbook.get_sample_data("grace_hopper.jpg")).mean(axis=-1)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12,5))
scat = ax1.scatter([0]*2000,[0]*2000,s=3, color='k')
im = ax2.imshow(img)
out = widgets.Output()
display(out)
#interact
def graph(
Spiral=True,
n=2000,
x1=50,
y1=50,
z1=50,
k1=300,
vlc=(0.1, 1, 0.01),
vuc=(0.1, 1, 0.01),
):
if Spiral == False:
x = 0
y = 0
else:
angle = np.linspace(x1, y1 * 1 * np.pi, n)
radius = np.linspace(z1, k1, n)
x = radius * np.cos(angle) + 150
y = radius * np.sin(angle) + 150
scat.set_offsets(np.c_[x, y])
# correctly scale the x and y limits
ax1.dataLim = scat.get_datalim(ax1.transData)
ax1.autoscale_view()
vu = np.quantile(img, vuc)
vl = np.quantile(img, vlc)
im.norm.vmin = vl
im.norm.vmax = vu
Approach 4 - mpl_interactions
Using set_offsets and equivalent set_data will be the most performant solution, but can also be tricky to figure out how to get it work and even trickier to remember. To make it easier I've creted a library (mpl-interactions) that automates the boilerplate of approach 3.
In addition to being easy and performant this has the advantage that you aren't responsible for updating the plots, only for returning the correct values. Which then has the ancillary benefit that now functions like spiral can be used in other parts of your code as they just return values rather than handle plotting.
The other advantage is that mpl-interactions can also create matplotlib widgets so this is the only approach that will also work outside of a notebook.
%matplotlib ipympl
import ipywidgets as widgets
import matplotlib.pyplot as plt
import numpy as np
import mpl_interactions.ipyplot as iplt
img = plt.imread(cbook.get_sample_data("grace_hopper.jpg")).mean(axis=-1)
# define the functions to be plotted
def spiral(Spiral=False, n=2000, x1=50, y1=50, z1=50, k1=300):
if Spiral == False:
x = 0
y = 0
return x, y
else:
angle = np.linspace(x1, y1 * 1 * np.pi, n)
radius = np.linspace(z1, k1, n)
x = radius * np.cos(angle) + 150
y = radius * np.sin(angle) + 150
return x, y
def vmin(vuc, vlc):
return np.quantile(img, vlc)
def vmax(vlc, vuc):
return np.quantile(img, vuc)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
controls = iplt.scatter(
spiral,
Spiral={(True, False)},
n=np.arange(1800, 2200),
x1=(25, 75),
y1=(25, 75),
z1=(25, 75),
k1=(200, 400),
parametric=True,
s=3,
c="black",
ax=ax1,
)
controls = iplt.imshow(
img,
vmin=vmin,
vmax=vmax,
vuc=(0.1, 1, 1000),
vlc=(0.1, 1, 1000),
controls=controls[None],
ax=ax2,
)
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:
Background
In Matplotlib, we can render the string using mathtext as a marker using $ ..... $ (Reference 1)
Question
Is there any way to enclose this text in a circular or rectangular box, or any different different shape? Similar to the registered symbol as shown here
I want to use this marker on a plot as shown below:
Text '$T$' is used in this plot, I want the text to be enclosed in a circle or rectangle.
Solution
As suggested in the comments of the answer, I have plotted a square marker of a bit larger size before the text marker. This resolved the issue.
The final figure is shown below:
Edit: Easiest way is to simply place patches to be the desired "frames" in the same location as the markers. Just make sure they have a lower zorder so that they don't cover the data points.
More sophisticated ways below:
You can make patches. Here is an example I used to make a custom question mark:
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.markers as m
fig, ax = plt.subplots()
lim = -5.8, 5.7
ax.set(xlim = lim, ylim = lim)
marker_obj = m.MarkerStyle('$?$') #Here you place your letter
path = marker_obj.get_path().transformed(marker_obj.get_transform())
path._vertices = np.array(path._vertices)*8 #To make it larger
patch = mpl.patches.PathPatch(path, facecolor="cornflowerblue", lw=2)
ax.add_patch(patch)
def translate_verts(patch, i=0, j=0, z=None):
patch._path._vertices = patch._path._vertices + [i, j]
def rescale_verts(patch, factor = 1):
patch._path._vertices = patch._path._vertices * factor
#translate_verts(patch, i=-0.7, j=-0.1)
circ = mpl.patches.Arc([0,0], 11, 11,
angle=0.0, theta1=0.0, theta2=360.0,
lw=10, facecolor = "cornflowerblue",
edgecolor = "black")
ax.add_patch(circ)#One of the rings around the questionmark
circ = mpl.patches.Arc([0,0], 10.5, 10.5,
angle=0.0, theta1=0.0, theta2=360.0,
lw=10, edgecolor = "cornflowerblue")
ax.add_patch(circ)#Another one of the rings around the question mark
circ = mpl.patches.Arc([0,0], 10, 10,
angle=0.0, theta1=0.0, theta2=360.0,
lw=10, edgecolor = "black")
ax.add_patch(circ)
if __name__ == "__main__":
ax.axis("off")
ax.set_position([0, 0, 1, 1])
fig.canvas.draw()
#plt.savefig("question.png", dpi=40)
plt.show()
Edit, second answer:
creating a custom patch made of other patches:
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import mpl_toolkits.mplot3d.art3d as art3d
class PlanetPatch(mpl.patches.Circle):
"""
This class combines many patches to make a custom patch
The best way to reproduce such a thing is to read the
source code for all patches you plan on combining.
Also make use of ratios as often as possible to maintain
proportionality between patches of different sizes"""
cz = 0
def __init__(self, xy, radius,
color = None, linewidth = 20,
edgecolor = "black", ringcolor = "white",
*args, **kwargs):
ratio = radius/6
mpl.patches.Circle.__init__(self, xy, radius,
linewidth = linewidth*ratio,
color = color,
zorder = PlanetPatch.cz,
*args, **kwargs)
self.set_edgecolor(edgecolor)
xy_ringcontour = np.array(xy)+[0, radius*-0.2/6]
self.xy_ringcontour = xy_ringcontour - np.array(xy)
self.ring_contour = mpl.patches.Arc(xy_ringcontour,
15*radius/6, 4*radius/6,
angle =10, theta1 = 165,
theta2 = 14.5,
fill = False,
linewidth = 65*linewidth*ratio/20,
zorder = 1+PlanetPatch.cz)
self.ring_inner = mpl.patches.Arc(xy_ringcontour,
15*radius/6, 4*radius/6,
angle = 10, theta1 = 165 ,
theta2 = 14.5,fill = False,
linewidth = 36*linewidth*ratio/20,
zorder = 2+PlanetPatch.cz)
self.top = mpl.patches.Wedge([0,0], radius, theta1 = 8,
theta2 = 192,
zorder=3+PlanetPatch.cz)
self.xy_init = xy
self.top._path._vertices=self.top._path._vertices+xy
self.ring_contour._edgecolor = self._edgecolor
self.ring_inner.set_edgecolor(ringcolor)
self.top._facecolor = self._facecolor
def add_to_ax(self, ax):
ax.add_patch(self)
ax.add_patch(self.ring_contour)
ax.add_patch(self.ring_inner)
ax.add_patch(self.top)
def translate(self, dx, dy):
self._center = self.center + [dx,dy]
self.ring_inner._center = self.ring_inner._center +[dx, dy]
self.ring_contour._center = self.ring_contour._center + [dx,dy]
self.top._path._vertices = self.top._path._vertices + [dx,dy]
def set_xy(self, new_xy):
"""As you can see all patches have different ways
to have their positions updated"""
new_xy = np.array(new_xy)
self._center = new_xy
self.ring_inner._center = self.xy_ringcontour + new_xy
self.ring_contour._center = self.xy_ringcontour + new_xy
self.top._path._vertices += new_xy - self.xy_init
fig = plt.figure(figsize=(6, 6))
ax = fig.add_subplot()
lim = -8.5, 8.6
ax.set(xlim = lim, ylim = lim,
facecolor = "black")
planets = []
colors = mpl.colors.cnames
colors = [c for c in colors]
for x in range(100):
xy = np.random.randint(-7, 7, 2)
r = np.random.randint(1, 15)/30
color = np.random.choice(colors)
planet = PlanetPatch(xy, r, linewidth = 20,
color = color,
ringcolor = np.random.choice(colors),
edgecolor = np.random.choice(colors))
planet.add_to_ax(ax)
planets.append(planet)
fig.canvas.draw()
#plt.savefig("planet.png", dpi=10)
plt.show()
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