plot_trisurface with custom color array - python

I basically want to "imshow" the pdf of a three-dimensional Dirichlet distribution on its support. Function simplex below computes regular points on that support, which are stored in the array sim. The array pdf holds a scalar density for each row in sim.
First thing I thought of was to use a triangulation. However, the color argument of plot_trisurface supports only one single color for all triangles. Setting cmap colors the triangles based on the z-coordinate values (See Fig. 1). Also plot_trisurface ignores the facecolors kwarg. What I want, however, is to color the surface based on pdf.
As a workaround I found, that I could interpolated the surface as 3d scatter plot. This generally gives the desired visualization, yet I ist clearly visible that it's a scatter plot; especially on the borders. (See Fig 2.)
Is there a way to plot the projection of the pdf onto the simplex?
import itertools
import matplotlib.pyplot as plt
import numpy as np
from scipy import stats
def simplex(n_vals):
base = np.linspace(0, 1, n_vals, endpoint=False)
coords = np.asarray(list(itertools.product(base, repeat=3)))
return coords[np.isclose(coords.sum(axis=-1), 1.0)]
sim = simplex(20)
pdf = stats.dirichlet([1.1, 1.5, 1.3]).pdf(sim.T)
fig1 = plt.figure()
ax1 = fig1.add_subplot(1, 2, 1, projection='3d', azim=20)
ax2 = fig1.add_subplot(1, 2, 2, projection='3d', azim=20)
ax1.plot_trisurf(x, y, z, color='k')
ax2.plot_trisurf(x, y, z, cmap='Spectral')
fig2 = plt.figure()
ax21 = fig2.add_subplot(projection='3d', azim=20)
ax21.scatter3D(*sim.T, s=50, alpha=.5, c=pdf, cmap='Spectral')

This is a way to do so by coloring each triangle in a triangulation object with the right color. Is this what you were looking for? The only thing is that each patch has a uniform color which make the patches somewhat visible.
# Setup is the same
import itertools
import matplotlib.pyplot as plt
from pylab import get_cmap
from matplotlib.tri import Triangulation, LinearTriInterpolator
import numpy as np
from scipy import stats
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
def simplex(n_vals):
base = np.linspace(0, 1, n_vals, endpoint=False)
coords = np.asarray(list(itertools.product(base, repeat=3)))
return coords[np.isclose(coords.sum(axis=-1), 1.0)]
sim = simplex(20)
pdf = stats.dirichlet([1.1, 1.5, 1.3]).pdf(sim.T)
# For shorter notation we define x, y and z:
x = sim[:, 0]
y = sim[:, 1]
z = sim[:, 2]
# Creating a triangulation object and using it to extract the actual triangles.
# Note if it is necessary that no patch will be vertical (i.e. along the z direction)
tri = Triangulation(x, y)
triangle_vertices = np.array([np.array([[x[T[0]], y[T[0]], z[T[0]]],
[x[T[1]], y[T[1]], z[T[1]]],
[x[T[2]], y[T[2]], z[T[2]]]]) for T in tri.triangles])
# Finding coordinate for the midpoints of each triangle.
# This will be used to extract the color
midpoints = np.average(triangle_vertices, axis = 1)
midx = midpoints[:, 0]
midy = midpoints[:, 1]
# Interpolating the pdf and using it with the selected cmap to produce the color RGB vector for each face.
# Some roundoff and normalization are needed
face_color_function = LinearTriInterpolator(tri, pdf)
face_color_index = face_color_function(midx, midy)
face_color_index[face_color_index < 0] = 0
face_color_index /= np.max(pdf)
cmap = get_cmap('Spectral')
# Creating the patches and plotting
collection = Poly3DCollection(triangle_vertices, facecolors=cmap(face_color_index), edgecolors=None)
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.add_collection(collection)
plt.show()
Obviously increasing the resolution would make the plot smoother.

This figure, complete with a colorbar,
was produced by the following script — the function map_colors, defined at the end of the script, could interest the general reader.
import numpy as np
import matplotlib.pyplot as plt
from scipy import stats
from itertools import product as Π
# the distribution that we want to study
dirichlet = stats.dirichlet([1.1, 1.5, 1.3])
# generate the "mesh"
N = 30 # no. of triangles along an edge
s = np.linspace(0, 1, N+1)
x, y, z = np.array([(x,y,1-x-y) for x,y in Π(s,s) if x+y<1+1E-6]).T
# plot as usual
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, projection='3d', azim=20)
p3dc = ax.plot_trisurf(x, y, z)
########## change the face colors ####################
mappable = map_colors(p3dc, dirichlet.pdf, 'Spectral')
# ####################################################
# possibly add a colormap
plt.colorbar(mappable, shrink=0.67, aspect=16.7)
# we are done
plt.show()
def map_colors(p3dc, func, cmap='viridis'):
"""
Color a tri-mesh according to a function evaluated in each barycentre.
p3dc: a Poly3DCollection, as returned e.g. by ax.plot_trisurf
func: a single-valued function of 3 arrays: x, y, z
cmap: a colormap NAME, as a string
Returns a ScalarMappable that can be used to instantiate a colorbar.
"""
from matplotlib.cm import ScalarMappable, get_cmap
from matplotlib.colors import Normalize
from numpy import array
# reconstruct the triangles from internal data
x, y, z, _ = p3dc._vec
slices = p3dc._segslices
triangles = array([array((x[s],y[s],z[s])).T for s in slices])
# compute the barycentres for each triangle
xb, yb, zb = triangles.mean(axis=1).T
# compute the function in the barycentres
values = func(xb, yb, zb)
# usual stuff
norm = Normalize()
colors = get_cmap(cmap)(norm(values))
# set the face colors of the Poly3DCollection
p3dc.set_fc(colors)
# if the caller wants a colorbar, they need this
return ScalarMappable(cmap=cmap, norm=norm)

Related

Create 3D Plot (not surface, scatter), where colour depends on z values

I want to create and save a number of sequential plots so I can then make an mp4 movie out of those plots. I want the color of the plot to depend on z (the value of the third axis):
The code I am using:
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator
import numpy as np
file_dir1 = r"C:\Users\files\final_files\B_6_sec\_read.csv"
specs23 = pd.read_csv(file_dir1, sep=',')
choose_file = specs23 # Choose file betwenn specs21, specs22,...
quant = 0 # Choose between 0,1,...,according to the following list
column = ['$\rho$', '$V_{x}$', '$V_{y}$', '$V_{z}$','$B_{x}$', '$B_{y}$','$B_{z}$','$Temperature$']
choose_column = choose_file[column[quant]]
resolution = 1024 # Specify resolution of grid
t_steps = int(len(specs23)/resolution) # Specify number of timesteps
fig, ax = plt.subplots(subplot_kw={"projection": "3d"},figsize=(15,10))
# Make data.
X = np.arange(0, resolution, 1)
Y = np.arange(0, int(len(specs23)/resolution),1)
X, Y = np.meshgrid(X, Y)
Z = choose_file[column[quant]].values
new_z = np.zeros((t_steps,resolution)) # Selected quantity as a function of x,t
### Plot figure ###
for i in range(0,int(len(choose_file)/resolution)):
zs = choose_column[i*resolution:resolution*(i+1)].values
new_z[i] = zs
for i in range(len(X)):
ax.plot(X[i], Y[i], new_z[i]) #%// color binded to "z" values
ax.zaxis.set_major_locator(LinearLocator(10))
# A StrMethodFormatter is used automatically
ax.zaxis.set_major_formatter('{x:.02f}')
plt.show()
What I am getting looks like this:
I would like to look it like this:
I have created the second plot using the LineCollection module. The problem is that it prints all the lines at once not allowing me to save each separately to create a movie.
You can find the dataframe I am using to create the figure here:
https://www.dropbox.com/s/idbeuhyxqfy9xvw/_read.csv?dl=0
The poster wants two things
lines with colors depending on z-values
animation of the lines over time
In order to achieve(1) one needs to cut up each line in separate segments and assign a color to each segment; in order to obtain a colorbar, we need to create a scalarmappable object that knows about the outer limits of the colors.
For achieving 2, one needs to either (a) save each frame of the animation and combine it after storing all the frames, or (b) leverage the animation module in matplotlib. I have used the latter in the example below and achieved the following:
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt, numpy as np
from mpl_toolkits.mplot3d.art3d import Line3DCollection
fig, ax = plt.subplots(subplot_kw = dict(projection = '3d'))
# generate data
x = np.linspace(-5, 5, 500)
y = np.linspace(-5, 5, 500)
z = np.exp(-(x - 2)**2)
# uggly
segs = np.array([[(x1,y2), (x2, y2), (z1, z2)] for x1, x2, y1, y2, z1, z2 in zip(x[:-1], x[1:], y[:-1], y[1:], z[:-1], z[1:])])
segs = np.moveaxis(segs, 1, 2)
# setup segments
# get bounds
bounds_min = segs.reshape(-1, 3).min(0)
bounds_max = segs.reshape(-1, 3).max(0)
# setup colorbar stuff
# get bounds of colors
norm = plt.cm.colors.Normalize(bounds_min[2], bounds_max[2])
cmap = plt.cm.plasma
# setup scalar mappable for colorbar
sm = plt.cm.ScalarMappable(norm, plt.cm.plasma)
# get average of segment
avg = segs.mean(1)[..., -1]
# get colors
colors = cmap(norm(avg))
# generate colors
lc = Line3DCollection(segs, norm = norm, cmap = cmap, colors = colors)
ax.add_collection(lc)
def update(idx):
segs[..., -1] = np.roll(segs[..., -1], idx)
lc.set_offsets(segs)
return lc
ax.set_xlim(bounds_min[0], bounds_max[0])
ax.set_ylim(bounds_min[1], bounds_max[1])
ax.set_zlim(bounds_min[2], bounds_max[2])
fig.colorbar(sm)
from matplotlib import animation
frames = np.linspace(0, 30, 10, 0).astype(int)
ani = animation.FuncAnimation(fig, update, frames = frames)
ani.save("./test_roll.gif", savefig_kwargs = dict(transparent = False))
fig.show()

Matplotlib: Set cmap in plot_surface to x and y-axes

How can I set the colormap in relation to the radius of the figure?
And how can I close the ends of the cylinder (on the element, not the top and bottom bases)?
My script:
import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from math import sin, cos, pi
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
h, w = 60,30
znew = np.random.randint(low=90, high=110, size=(60,30))
theta = np.linspace(0,2*pi, h)
Z = np.linspace(0,1,w)
Z,theta = np.meshgrid(Z, theta)
R = 1
X = (R*np.cos(theta))*znew
Y = (R*np.sin(theta))*znew
ax1 = ax.plot_surface(X,Y,Z,linewidth = 0, cmap="coolwarm",
vmin= 80,vmax=130, shade = True, alpha = 0.75)
fig.colorbar(ax1, shrink=0.9, aspect=5)
plt.show()
First you need to use the facecolors keyword argument of plot_surface to draw your surface with arbitrary (non-Z-based) colours. You have to pass an explicit RGBA colour four each point, which means we need to sample a colormap object with the keys given by the radius at every point. Finally, this will break the mappable property of the resulting surface, so we will have to construct the colorbar by manually telling it to use our radii for colours:
import numpy as np
from matplotlib import pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import Normalize
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
h, w = 60,30
#znew = np.random.randint(low=90, high=110, size=(h,w))
theta = np.linspace(0,2*np.pi, h)
Z = np.linspace(0,1,w)
Z,theta = np.meshgrid(Z, theta)
znew = 100 + 10*np.cos(theta/2)*np.cos(2*Z*np.pi)
R = 1
X = (R*np.cos(theta))*znew
Y = (R*np.sin(theta))*znew
true_radius = np.sqrt(X**2 + Y**2)
norm = Normalize()
colors = norm(true_radius) # auto-adjust true radius into [0,1] for color mapping
cmap = cm.get_cmap("coolwarm")
ax.plot_surface(X, Y, Z, linewidth=0, facecolors=cmap(colors), shade=True, alpha=0.75)
# the surface is not mappable, we need to handle the colorbar manually
mappable = cm.ScalarMappable(cmap=cmap)
mappable.set_array(colors)
fig.colorbar(mappable, shrink=0.9, aspect=5)
plt.show()
Note that I changed the radii to something smooth for a less chaotic-looking result. The true_radius arary contains the actual radii in data units, which after normalization becomes colors (essentially colors = (true_radius - true_radius.min())/true_radius.ptp()).
The result:
Finally, note that I generated the radii such that the cylinder doesn't close seamlessly. This mimicks your random example input. There's nothing you can do about this as long as the radii are not 2π-periodic in theta. This has nothing to do with visualization, this is geometry.

Python/Matplotlib: 2d random walk with kde joint density contour in a 3d plot

I'm struggling with creating a quite complex 3d figure in python, specifically using iPython notebook. I can partition the content of the graph into two sections:
The (x,y) plane: Here a two-dimensional random walk is bobbing around, let's call it G(). I would like to plot part of this trajectory on the (x,y) plane. Say, 10% of all the data points of G(). As G() bobs around, it visits some (x,y) pairs more frequently than others. I would like to estimate this density of G() using a kernel estimation approach and draw it as contour lines on the (x,y) plane.
The (z) plane: Here, I would like to draw a mesh or (transparent) surface plot of the information theoretical surprise of a bivariate normal. Surprise is simply -log(p(i)) or the negative (base 2) logarithm of outcome i. Given the bivariate normal, each (x,y) pair has some probability p(x,y) and the surprise of this is simply -log(p(x,y)).
Essentially these two graphs are independent. Assume the interval of the random walk G() is [xmin,xmax],[ymin,ymax] and of size N. The bivariate normal in the z-plane should be drawn from the same interval, such that for each (x,y) pair in the random walk, I can draw a (dashed) line from some subset of the random walk n < N to the bivariate normal. Assume that G(10) = (5,5) then I would like to draw a dashed line from (5,5) up the Z-axes, until it hits the bivariate normal.
So far, I've managed to plot G() in a 3-d space, and estimate the density f(X,Y) using scipy.stats.gaussian_kde. In another (2d) graph, I have the sort of contour lines I want. What I don't have, is the contour lines in the 3d-plot using the estimated KDE density. I also don't have the bivariate normal plot, or the projection of a few random points from the random walk, to the surface of the bivariate normal. I've added a hand drawn figure, which might ease intuition (ignore the label on the z-axis and the fact that there is no mesh.. difficult to draw!)
Any input, even just partial, such as how to draw the contour lines in the (x,y) plane of the 3d graph, or a mesh of a bivariate normal would be much appreciated.
Thanks!
import matplotlib as mpl
import matplotlib.pyplot as plt
import random
import numpy as np
import seaborn as sns
import scipy
from mpl_toolkits.mplot3d import Axes3D
%matplotlib inline
def randomwalk():
mpl.rcParams['legend.fontsize'] = 10
xyz = []
cur = [0, 0]
for _ in range(400):
axis = random.randrange(0, 2)
cur[axis] += random.choice([-1, 1])
xyz.append(cur[:])
x, y = zip(*xyz)
data = np.vstack([x,y])
kde = scipy.stats.gaussian_kde(data)
density = kde(data)
fig1 = plt.figure()
ax = fig1.gca(projection='3d')
ax.plot(x, y, label='Random walk')
sns.kdeplot(data[0,:], data[1,:], 0)
ax.scatter(x[-1], y[-1], c='b', marker='o') # End point
ax.legend()
fig2 = plt.figure()
sns.kdeplot(data[0,:], data[1,:])
Calling randomwalk() initialises and plots this:
Edit #1:
Made some progress, actually the only thing I need is to restrict the height of the dashed vertical lines to the bivariate. Any ideas?
import matplotlib as mpl
import matplotlib.pyplot as plt
import random
import numpy as np
import seaborn as sns
import scipy
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.mlab import bivariate_normal
%matplotlib inline
# Data for random walk
def randomwalk():
mpl.rcParams['legend.fontsize'] = 10
xyz = []
cur = [0, 0]
for _ in range(40):
axis = random.randrange(0, 2)
cur[axis] += random.choice([-1, 1])
xyz.append(cur[:])
# Get density
x, y = zip(*xyz)
data = np.vstack([x,y])
kde = scipy.stats.gaussian_kde(data)
density = kde(data)
# Data for bivariate gaussian
a = np.linspace(-7.5, 7.5, 20)
b = a
X,Y = np.meshgrid(a, b)
Z = bivariate_normal(X, Y)
surprise_Z = -np.log(Z)
# Get random points from walker and plot up z-axis to the gaussian
M = data[:,np.random.choice(20,5)].T
# Plot figure
fig = plt.figure(figsize=(10, 7))
ax = fig.gca(projection='3d')
ax.plot(x, y, 'grey', label='Random walk') # Walker
ax.scatter(x[-1], y[-1], c='k', marker='o') # End point
ax.legend()
surf = ax.plot_surface(X, Y, surprise_Z, rstride=1, cstride=1,
cmap = plt.cm.gist_heat_r, alpha=0.1, linewidth=0.1)
#fig.colorbar(surf, shrink=0.5, aspect=7, cmap=plt.cm.gray_r)
for i in range(5):
ax.plot([M[i,0], M[i,0]],[M[i,1], M[i,1]], [0,10],'k--',alpha=0.8, linewidth=0.5)
ax.set_zlim(0, 50)
ax.set_xlim(-10, 10)
ax.set_ylim(-10, 10)
Final code,
import matplotlib as mpl
import matplotlib.pyplot as plt
import random
import numpy as np
import seaborn as sns
import scipy
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.mlab import bivariate_normal
%matplotlib inline
# Data for random walk
def randomwalk():
mpl.rcParams['legend.fontsize'] = 10
xyz = []
cur = [0, 0]
for _ in range(50):
axis = random.randrange(0, 2)
cur[axis] += random.choice([-1, 1])
xyz.append(cur[:])
# Get density
x, y = zip(*xyz)
data = np.vstack([x,y])
kde = scipy.stats.gaussian_kde(data)
density = kde(data)
# Data for bivariate gaussian
a = np.linspace(-7.5, 7.5, 100)
b = a
X,Y = np.meshgrid(a, b)
Z = bivariate_normal(X, Y)
surprise_Z = -np.log(Z)
# Get random points from walker and plot up z-axis to the gaussian
M = data[:,np.random.choice(50,10)].T
# Plot figure
fig = plt.figure(figsize=(10, 7))
ax = fig.gca(projection='3d')
ax.plot(x, y, 'grey', label='Random walk') # Walker
ax.legend()
surf = ax.plot_surface(X, Y, surprise_Z, rstride=1, cstride=1,
cmap = plt.cm.gist_heat_r, alpha=0.1, linewidth=0.1)
#fig.colorbar(surf, shrink=0.5, aspect=7, cmap=plt.cm.gray_r)
for i in range(10):
x = [M[i,0], M[i,0]]
y = [M[i,1], M[i,1]]
z = [0,-np.log(bivariate_normal(M[i,0],M[i,1]))]
ax.plot(x,y,z,'k--',alpha=0.8, linewidth=0.5)
ax.scatter(x, y, z, c='k', marker='o')

matplotlib scatterplot: adding 4th dimension by the marker shape

I would like to add a fourth dimension to the scatter plot by defining the ellipticity of the markers depending on a variable. Is that possible somehow ?
EDIT:
I would like to avoid a 3D-plot. In my opinion these plots are usually not very informative.
You can place Ellipse patches directly onto your axes, as demonstrated in this matplotlib example. To adapt it to use eccentricity as your "third dimension") keeping the marker area constant:
from pylab import figure, show, rand
from matplotlib.patches import Ellipse
import numpy as np
import matplotlib.pyplot as plt
N = 25
# ellipse centers
xy = np.random.rand(N, 2)*10
# ellipse eccentrities
eccs = np.random.rand(N) * 0.8 + 0.1
fig = plt.figure()
ax = fig.add_subplot(111, aspect='equal')
A = 0.1
for pos, e in zip(xy, eccs):
# semi-minor, semi-major axes, b and a:
b = np.sqrt(A/np.pi * np.sqrt(1-e**2))
a = A / np.pi / b
ellipse = Ellipse(xy=pos, width=2*a, height=2*b)
ax.add_artist(ellipse)
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
show()
Of course, you need to scale your marker area to your x-, y- values in this case.
You can use colorbar as the 4th dimension to your 3D plot. One example is as shown below:
import matplotlib.cm as cmx
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import matplotlib
import numpy as np
def scatter3d(x,y,z, cs, colorsMap='jet'):
cm = plt.get_cmap(colorsMap)
cNorm = matplotlib.colors.Normalize(vmin=min(cs), vmax=max(cs))
scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=cm)
fig = plt.figure()
ax = Axes3D(fig)
ax.scatter(x, y, z, c=scalarMap.to_rgba(cs))
scalarMap.set_array(cs)
fig.colorbar(scalarMap,label='Test')
plt.show()
x = np.random.uniform(0,1,50)
y = np.random.uniform(0,1,50)
z = np.random.uniform(0,1,50)
so scatter3D(x,y,z,x+y) produces:
with x+y being the 4th dimension shown in color. You can add your calculated ellipticity depending on your specific variable instead of x+y to get what you want.
To change the ellipticity of the markers you will have to create them manually as such a feature is not implemented yet. However, I believe you can show 4 dimensions with a 2D scatter plot by using color and size as additional dimensions. You will have to take care of the scaling from data to marker size yourself. I added a simple function to handle that in the example below:
import matplotlib.pyplot as plt
import numpy as np
data = np.random.rand(60,4)
def scale_size(data, data_min=None, data_max=None, size_min=10, size_max=60):
# if the data limits are set to None we will just infer them from the data
if data_min is None:
data_min = data.min()
if data_max is None:
data_max = data.max()
size_range = size_max - size_min
data_range = data_max - data_min
return ((data - data_min) * size_range / data_range) + size_min
plt.scatter(data[:,0], data[:,1], c=data[:,2], s=scale_size(data[:,3]))
plt.colorbar()
plt.show()
Result:

Is it possible to change line color in a plot if exceeds a specific range?

Is it possible to change the line color in a plot when values exceeds a certain y value?
Example:
import numpy as np
import matplotlib.pyplot as plt
a = np.array([1,2,17,20,16,3,5,4])
plt.plt(a)
This one gives the following:
I want to visualise the values that exceeds y=15. Something like the following figure:
Or something like this(with cycle linestyle)::
Is it possible?
Define a helper function (this a bare-bones one, more bells and whistles can be added). This code is a slight refactoring of this example from the documentation.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import ListedColormap, BoundaryNorm
def threshold_plot(ax, x, y, threshv, color, overcolor):
"""
Helper function to plot points above a threshold in a different color
Parameters
----------
ax : Axes
Axes to plot to
x, y : array
The x and y values
threshv : float
Plot using overcolor above this value
color : color
The color to use for the lower values
overcolor: color
The color to use for values over threshv
"""
# Create a colormap for red, green and blue and a norm to color
# f' < -0.5 red, f' > 0.5 blue, and the rest green
cmap = ListedColormap([color, overcolor])
norm = BoundaryNorm([np.min(y), threshv, np.max(y)], cmap.N)
# Create a set of line segments so that we can color them individually
# This creates the points as a N x 1 x 2 array so that we can stack points
# together easily to get the segments. The segments array for line collection
# needs to be numlines x points per line x 2 (x and y)
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
# Create the line collection object, setting the colormapping parameters.
# Have to set the actual values used for colormapping separately.
lc = LineCollection(segments, cmap=cmap, norm=norm)
lc.set_array(y)
ax.add_collection(lc)
ax.set_xlim(np.min(x), np.max(x))
ax.set_ylim(np.min(y)*1.1, np.max(y)*1.1)
return lc
Example of usage
fig, ax = plt.subplots()
x = np.linspace(0, 3 * np.pi, 500)
y = np.sin(x)
lc = threshold_plot(ax, x, y, .75, 'k', 'r')
ax.axhline(.75, color='k', ls='--')
lc.set_linewidth(3)
and the output
If you want just the markers to change color, use the same norm and cmap and pass them to scatter as
cmap = ListedColormap([color, overcolor])
norm = BoundaryNorm([np.min(y), threshv, np.max(y)], cmap.N)
sc = ax.scatter(x, y, c=c, norm=norm, cmap=cmap)
Unfortunately, matplotlib doesn't have an easy option to change the color of only part of a line. We will have to write the logic ourselves. The trick is to cut the line up into a collection of line segments, then assign a color to each of them, and then plot them.
from matplotlib import pyplot as plt
from matplotlib.collections import LineCollection
import numpy as np
# The x and y data to plot
y = np.array([1,2,17,20,16,3,5,4])
x = np.arange(len(y))
# Threshold above which the line should be red
threshold = 15
# Create line segments: 1--2, 2--17, 17--20, 20--16, 16--3, etc.
segments_x = np.r_[x[0], x[1:-1].repeat(2), x[-1]].reshape(-1, 2)
segments_y = np.r_[y[0], y[1:-1].repeat(2), y[-1]].reshape(-1, 2)
# Assign colors to the line segments
linecolors = ['red' if y_[0] > threshold and y_[1] > threshold else 'blue'
for y_ in segments_y]
# Stamp x,y coordinates of the segments into the proper format for the
# LineCollection
segments = [zip(x_, y_) for x_, y_ in zip(segments_x, segments_y)]
# Create figure
plt.figure()
ax = plt.axes()
# Add a collection of lines
ax.add_collection(LineCollection(segments, colors=linecolors))
# Set x and y limits... sadly this is not done automatically for line
# collections
ax.set_xlim(0, 8)
ax.set_ylim(0, 21)
Your second option is much easier. We first draw the line and then add the markers as a scatterplot on top of it:
from matplotlib import pyplot as plt
import numpy as np
# The x and y data to plot
y = np.array([1,2,17,20,16,3,5,4])
x = np.arange(len(y))
# Threshold above which the markers should be red
threshold = 15
# Create figure
plt.figure()
# Plot the line
plt.plot(x, y, color='blue')
# Add below threshold markers
below_threshold = y < threshold
plt.scatter(x[below_threshold], y[below_threshold], color='green')
# Add above threshold markers
above_threshold = np.logical_not(below_threshold)
plt.scatter(x[above_threshold], y[above_threshold], color='red')
Basically #RaJa provides the solution, but I think that you can do the same without loading an additional package (pandas), by using masked arrays in numpy:
import numpy as np
import matplotlib.pyplot as plt
a = np.array([1,2,17,20,16,3,5,4])
# use a masked array to suppress the values that are too low
a_masked = np.ma.masked_less_equal(a, 15)
# plot the full line
plt.plot(a, 'k')
# plot only the large values
plt.plot(a_masked, 'r', linewidth=2)
# add the threshold value (optional)
plt.axhline(15, color='k', linestyle='--')
plt.show()
Result:
I don't know wether there is a built-in function in matplolib. But you could convert your numpy array into a pandas series and then use the plot function in combination with boolean selection/masking.
import numpy as np
import pandas as pd
a = np.array([1,2,17,20,16,3,5,4])
aPandas = pd.Series(a)
aPandas.plot()
aPandas[aPandas > 15].plot(color = 'red')

Categories