plot_forcefield.py 文件源码

def porestreamlines(polygon=None, rx=10., ry=10., Nx=100, Ny=100,
                    maxvalue=None, **fields):
    "streamlines plot of vector field around nanopore"

    # interpolate on regular mesh symmetric w.r.t. center axis
    mesh2D = nanopores.RectangleMesh([-rx-0.1,-ry-0.1], [rx+0.1,ry+0.1], Nx, Ny)
    fields2 = nanopores.convert2D(mesh2D, *(fields.values()))

    # prepare polygon and copy to left half
    settings = dict(closed=True, facecolor="#eeeeee", linewidth=3.,
                    edgecolor="black")
    if polygon is not None:
        polygon = np.array(polygon)
        polygon_m = np.column_stack([-polygon[:,0], polygon[:,1]])

    # prepare plots
    Ny += 1
    Nx += 1
    Y, X = np.mgrid[-ry:ry:Ny*1j, -rx:rx:Nx*1j]
    U = np.zeros((Ny,Nx))
    V = np.zeros((Ny,Nx))
    formt = matplotlib.ticker.FuncFormatter(fmt)
    ticks = [0] + [10**n for n in range(-16, -9)]

    # determine uniform color range from fields (maybe round to nearest 10-power)
    if maxvalue is None:
        maxvalue = max(dolfin.norm(F.vector(), "linf") for F in fields2)
        #maxvalue = 10**int(np.log10(maxvalue))

    for i, F in enumerate(fields2):
        Fstr = fields.keys()[i]
        fig, ax = plt.subplots(figsize=(5, 4.5), num=Fstr)

        # fill array with function values
        for y in range(Ny):
            for x in range(Nx):
                f = F(X[y][x], Y[y][x])
                U[y][x] = f[0]
                V[y][x] = f[1]

        # streamplot with logarithmic scale
        strength = np.sqrt(U*U+V*V)
        norm = matplotlib.colors.SymLogNorm(linthresh=ticks[1], linscale=1.0,
                                            vmin=0., vmax=maxvalue)
        strm = plt.streamplot(X, Y, U, V, arrowsize=1.5, linewidth=1.5, density=1.5,
                              cmap=cm.viridis, color=strength, norm=norm)
        plt.colorbar(strm.lines, ticks=ticks, format=formt)
        plt.xlabel('x [nm]') #, fontsize=20)
        plt.ylabel('z [nm]') #, fontsize=20)

        # plot pore polygon on top
        if polygon is not None:
            patch = patches.Polygon(polygon, **settings)
            patchm = patches.Polygon(polygon_m, **settings)
            patch.set_zorder(10)
            patchm.set_zorder(10)
            ax.add_patch(patch)
            ax.add_patch(patchm)