matplotlib streamplot上的箭头数量

时间:2017-03-31 23:09:40

标签: python matplotlib

无论如何增加matplotlib streamplot上的箭头数量?现在,似乎三个每个流线只有一个箭头,如果我想更改为x / y轴限制以放大数据,这就是一个问题。

3 个答案:

答案 0 :(得分:2)

我不确定只增加箭头的数量 - 但您可以使用 streamplot 功能中的密度参数来增加流线的密度,此处和#39; s文档:

*density* : float or 2-tuple
    Controls the closeness of streamlines. When `density = 1`, the domain
    is divided into a 30x30 grid---*density* linearly scales this grid.
    Each cell in the grid can have, at most, one traversing streamline.
    For different densities in each direction, use [density_x, density_y].

以下是一个例子:

import matplotlib.pyplot as plt
import numpy as np

x = np.arange(0,20,1)
y = np.arange(0,20,1)

u=np.random.random((x.shape[0], y.shape[0]))
v=np.random.random((x.shape[0], y.shape[0]))


fig, ax = plt.subplots(2,2)

ax[0,0].streamplot(x,y,u,v,density=1)
ax[0,0].set_title('Original')

ax[0,1].streamplot(x,y,u,v,density=4)
ax[0,1].set_xlim(5,10)
ax[0,1].set_ylim(5,10)
ax[0,1].set_title('Zoomed, higher density')

ax[1,1].streamplot(x,y,u,v,density=1)
ax[1,1].set_xlim(5,10)
ax[1,1].set_ylim(5,10)
ax[1,1].set_title('Zoomed, same density')

ax[1,0].streamplot(x,y,u,v,density=4)
ax[1,0].set_title('Original, higher density')

fig.show()

enter image description here

答案 1 :(得分:1)

我找到了一种自定义流线图中箭头数量的方法。

想法是分别绘制流线和箭头:

  • plt.streamplot返回具有两个属性的stream_container:linesarrowslines包含可用于重建无箭头流线的线段。
  • plt.quiver可用于绘制梯度场。如果缩放比例正确,箭头的长度可以忽略不计,只留下箭头。

因此,我们只需要使用线段定义箭头的位置并将其传递到plt.quiver

这是一个玩具示例:

import matplotlib.pyplot as plt
from matplotlib import collections as mc
import numpy as np

# get line segments
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
sp = ax.streamplot(x, y, u, v, start_points=start_points, density=10)
seg = sps.lines.get_segments()  # seg is a list of (2, 2) numpy arrays
lc = mc.LineCollection(seg, ...)

# define arrows
# here I define one arrow every 50 segments
# you could also select segs based on some criterion, e.g. intersect with certain lines
period = 50
arrow_x = np.array([seg[i][0, 0] for i in range(0, len(seg), period)])
arrow_y = np.array([seg[i][0, 1] for i in range(0, len(seg), period)])
arrow_dx = np.array([seg[i][1, 0] - seg[i][0, 0] for i in range(0, len(seg), period)])
arrow_dy = np.array([seg[i][1, 1] - seg[i][0, 1] for i in range(0, len(seg), period)])

# plot the final streamline
fig = plt.figure(figsize=(12.8, 10.8))
ax = fig.add_subplot(1, 1, 1)
ax.add_collection(lc)
ax.autoscale()
ax.quiver(
    arrow_x, arrow_y, arrow_dx, arrow_dy, angles='xy',  # arrow position
    scale=0.2, scale_units='inches', units='y', minshaft=0,  # arrow scaling
    headwidth=6, headlength=10, headaxislength=9)  # arrow style
fig.show()

有多种方法可以缩放箭头,使箭头的长度看起来为零。

答案 2 :(得分:1)

基于@Richard_wth的答案,我编写了一个函数来提供对流图上箭头位置的控制。您可以为每个流线选择n箭头,或选择使箭头在流线上等距分布。

首先,您执行常规的streamplot,直到对流线的位置和数量感到满意为止。您保留返回的参数sp。例如:

sp = ax.streamplot(x,y,u,v,arrowstyle='-',density=10)

这里重要的是要有arrowstyle='-'以便不显示箭头。

然后,您可以调用函数streamQuiver(下面提供)来控制每个流线上的箭头。如果您希望每个流线3个箭头:

streamQuiver(ax, sp, n=3, ...)

如果要每1.5个曲线长度进行精简:

streamQuiver(ax, sp, spacing=1.5, ...)

其中...是将传递给quiver的选项。 函数streamQuiver可能不是完全防弹的,在某些情况下可能需要一些其他处理。它依赖于4个子功能:

  • curve_coord获取沿路径的曲线长度
  • curve extract提取路径上的等距点
  • seg_to_lines将片段从流图转换为连续线。可能会有更好的方法!
  • lines_to_arrows:这是提取每行箭头的主要功能

在此示例中,箭头在每条流线上的等距点处。

import numpy as np
import matplotlib.pyplot as plt

def streamQuiver(ax,sp,*args,spacing=None,n=5,**kwargs):
    """ Plot arrows from streamplot data  
    The number of arrows per streamline is controlled either by `spacing` or by `n`.
    See `lines_to_arrows`.
    """
    def curve_coord(line=None):
        """ return curvilinear coordinate """
        x=line[:,0]
        y=line[:,1]
        s     = np.zeros(x.shape)
        s[1:] = np.sqrt((x[1:]-x[0:-1])**2+ (y[1:]-y[0:-1])**2)
        s     = np.cumsum(s)                                  
        return s

    def curve_extract(line,spacing,offset=None):
        """ Extract points at equidistant space along a curve"""
        x=line[:,0]
        y=line[:,1]
        if offset is None:
            offset=spacing/2
        # Computing curvilinear length
        s = curve_coord(line)
        offset=np.mod(offset,s[-1]) # making sure we always get one point
        # New (equidistant) curvilinear coordinate
        sExtract=np.arange(offset,s[-1],spacing)
        # Interpolating based on new curvilinear coordinate
        xx=np.interp(sExtract,s,x);
        yy=np.interp(sExtract,s,y);
        return np.array([xx,yy]).T

    def seg_to_lines(seg):
        """ Convert a list of segments to a list of lines """ 
        def extract_continuous(i):
            x=[]
            y=[]
            # Special case, we have only 1 segment remaining:
            if i==len(seg)-1:
                x.append(seg[i][0,0])
                y.append(seg[i][0,1])
                x.append(seg[i][1,0])
                y.append(seg[i][1,1])
                return i,x,y
            # Looping on continuous segment
            while i<len(seg)-1:
                # Adding our start point
                x.append(seg[i][0,0])
                y.append(seg[i][0,1])
                # Checking whether next segment continues our line
                Continuous= all(seg[i][1,:]==seg[i+1][0,:])
                if not Continuous:
                    # We add our end point then
                    x.append(seg[i][1,0])
                    y.append(seg[i][1,1])
                    break
                elif i==len(seg)-2:
                    # we add the last segment
                    x.append(seg[i+1][0,0])
                    y.append(seg[i+1][0,1])
                    x.append(seg[i+1][1,0])
                    y.append(seg[i+1][1,1])
                i=i+1
            return i,x,y
        lines=[]
        i=0
        while i<len(seg):
            iEnd,x,y=extract_continuous(i)
            lines.append(np.array( [x,y] ).T)
            i=iEnd+1
        return lines

    def lines_to_arrows(lines,n=5,spacing=None,normalize=True):
        """ Extract "streamlines" arrows from a set of lines 
        Either: `n` arrows per line
            or an arrow every `spacing` distance
        If `normalize` is true, the arrows have a unit length
        """
        if spacing is None:
            # if n is provided we estimate the spacing based on each curve lenght)
            spacing = [ curve_coord(l)[-1]/n for l in lines]
        try:
            len(spacing)
        except:
            spacing=[spacing]*len(lines)

        lines_s=[curve_extract(l,spacing=sp,offset=sp/2)         for l,sp in zip(lines,spacing)]
        lines_e=[curve_extract(l,spacing=sp,offset=sp/2+0.01*sp) for l,sp in zip(lines,spacing)]
        arrow_x  = [l[i,0] for l in lines_s for i in range(len(l))]
        arrow_y  = [l[i,1] for l in lines_s for i in range(len(l))]
        arrow_dx = [le[i,0]-ls[i,0] for ls,le in zip(lines_s,lines_e) for i in range(len(ls))]
        arrow_dy = [le[i,1]-ls[i,1] for ls,le in zip(lines_s,lines_e) for i in range(len(ls))]

        if normalize:
            dn = [ np.sqrt(ddx**2 + ddy**2) for ddx,ddy in zip(arrow_dx,arrow_dy)]
            arrow_dx = [ddx/ddn for ddx,ddn in zip(arrow_dx,dn)] 
            arrow_dy = [ddy/ddn for ddy,ddn in zip(arrow_dy,dn)] 
        return  arrow_x,arrow_y,arrow_dx,arrow_dy 

    # --- Main body of streamQuiver
    # Extracting lines
    seg   = sp.lines.get_segments() # list of (2, 2) numpy arrays
    lines = seg_to_lines(seg)       # list of (N,2) numpy arrays
    # Convert lines to arrows
    ar_x, ar_y, ar_dx, ar_dy = lines_to_arrows(lines,spacing=spacing,n=n,normalize=True)
    # Plot arrows
    qv=ax.quiver(ar_x, ar_y, ar_dx, ar_dy, *args, angles='xy', **kwargs)
    return qv

# --- Example
x = np.linspace(-1,1,100)
y = np.linspace(-1,1,100)
X,Y=np.meshgrid(x,y)
u = -np.sin(np.arctan2(Y,X))
v =  np.cos(np.arctan2(Y,X))

xseed=np.linspace(0.1,1,4)

fig=plt.figure()
ax=fig.add_subplot(111)
sp = ax.streamplot(x,y,u,v,color='k',arrowstyle='-',start_points=np.array([xseed,xseed*0]).T,density=30)
qv = streamQuiver(ax,sp,spacing=0.5, scale=60)
plt.show()

Example where each arrow is at a given distance along the streamline