Python plot_dofs._get_axes函数代码示例

def mark_subdomains(ax, cmesh, cell_tasks,
                    size=None, icolor=0, alpha=1.0, mask=False):
    """
    Mark cells of subdomains corresponding to each task by a different color.
    Plots nothing in 3D.
    """
    if size is None:
        size = cell_tasks.max() + 1

    coors = cmesh.coors
    dim = cmesh.dim

    ax = pd._get_axes(ax, dim)

    if dim == 3:
        return ax

    conn = cmesh.get_conn(dim, 0)

    color = nm.zeros(4)
    color[-1] = alpha
    for ic, vertices in enumerate(conn.indices.reshape((conn.num, -1))):
        cv = coors[vertices]
        if dim == 2:
            if mask:
                cc = cv.mean(0)
                cv = cc + 0.3 * (cv - cc)

            if not mask or cell_tasks[ic] > 0:
                color[icolor] = (float(cell_tasks[ic]) + 1) / size
                ax.fill(*cv.T, color=color)

    return ax

def plot_control_mesh(ax, control_points, label=False):
    """
    Plot the control mesh of a NURBS given by its control points.
    """
    dim = control_points.shape[-1]
    ax = _get_axes(ax, dim)

    shape = control_points.shape

    conn, desc = get_tensor_product_conn(nm.array(shape[:-1]))
    gel = GeometryElement(desc)

    coors = control_points.reshape((-1, dim))

    ax = pd.plot_mesh(ax, coors, conn, gel.edges)
    pd.plot_points(ax, coors)

    if label:
        for ii, cc in enumerate(coors):
            if dim == 3:
                cx, cy, cz = cc
                ax.text(cx, cy, cz, '%d' % ii,
                        color='g', fontsize=12, weight='bold')

            else:
                cx, cy = cc
                ax.text(cx, cy, '%d' % ii,
                        color='g', fontsize=12, weight='bold')

    return ax

def label_local_entities(ax, cmesh, edim, color='b', fontsize=10, **kwargs):
    """
    Label mesh topology entities using cell-local ids.
    """
    coors = cmesh.get_centroids(edim)
    coors = _to2d(coors)
    dim = cmesh.dim
    centres = cmesh.get_centroids(dim)

    cmesh.setup_connectivity(dim, edim)
    conn = cmesh.get_conn(dim, edim)
    off = conn.offsets

    ax = _get_axes(ax, dim)

    eps = 0.015 * fontsize
    oeps = 1.0 - eps
    for ii in range(conn.num):
        for ic, ie in enumerate(conn.indices[off[ii]:off[ii+1]]):
            # Shift labels towards the cell centre.
            cc = oeps * coors[ie] + eps * centres[ii]
            ax.text(*cc.T, s=ic, color=color, fontsize=fontsize,
                    horizontalalignment='center', verticalalignment='center',
                    **kwargs)

    return ax

def plot_edges(ax, gel, length, show=False):
    """
    Plot edges of a geometry element as numbered arrows.
    """
    dim = gel.dim
    ax = _get_axes(ax, dim)

    l2 = 0.5 * length
    for ii, edge in enumerate(gel.edges):
        cc = gel.coors[edge]
        centre = 0.5 * cc.sum(axis=0)

        vdir = (cc - centre)
        normalize_vectors(vdir)

        cc = l2 * vdir + centre
        draw_arrow(ax, cc, length=0.3*length, linewidth=3, color='b')

        if dim == 3:
            cx, cy, cz = centre
            ax.text(cx, cy, cz, ii,
                    color='b', fontsize=10, weight='light')

        else:
            cx, cy = centre
            ax.text(cx, cy, ii,
                    color='b', fontsize=10, weight='light')

    return ax

def plot_quadrature(ax, geometry, order, boundary=False,
                    min_radius=10, max_radius=50,
                    show_colorbar=False, show_labels=False):
    """
    Plot quadrature points for the given geometry and integration order.

    The points are plotted as circles/spheres with radii given by quadrature
    weights - the weights are mapped to [`min_radius`, `max_radius`] interval.
    """
    if not boundary:
        gel, coors, weights = _get_qp(geometry, order)

    else:
        gel, coors, weights = _get_bqp(geometry, order)

    dim = coors.shape[1]

    ax = _get_axes(ax, dim)

    plot_geometry(ax, gel)
    plot_weighted_points(ax, coors, weights,
                         min_radius=min_radius, max_radius=max_radius,
                         show_colorbar=show_colorbar)
    if show_labels:
        label_points(ax, coors)

    return ax, coors, weights

def label_local_entities(ax, cmesh, edim, color='b', fontsize=10, show=False):
    """
    Label mesh topology entities using cell-local ids.
    """
    coors = cmesh.get_centroids(edim)
    dim = cmesh.dim
    centres = cmesh.get_centroids(dim)

    conn = cmesh.get_conn(dim, edim)
    off = conn.offsets

    ax = _get_axes(ax, dim)

    eps = 0.1
    oeps = 1.0 - eps
    for ii in xrange(conn.num):
        for ic, ie in enumerate(conn.indices[off[ii]:off[ii+1]]):
            # Shift labels towards the cell centre.
            cc = oeps * coors[ie] + eps * centres[ii]
            if dim == 3:
                ax.text(cc[0], cc[1], cc[2], ic,
                        color=color, fontsize=fontsize)

            else:
                ax.text(cc[0], cc[1], ic,
                        color=color, fontsize=fontsize)

    if show:
        plt.show()

    return ax

def plot_nurbs_basis_1d(ax, nurbs, n_points=100, x_axis="parametric", legend=False):
    """
    Plot a 1D NURBS basis.
    """
    ax = _get_axes(ax, 2)

    ga = nurbs.greville()[0]

    n_fun = nurbs.weights.shape[0]
    line = nm.linspace(ga[0], ga[-1], n_points)
    for ii in range(n_fun):
        field = nm.zeros(n_fun)
        field[ii] = 1.0

        vals = nurbs.evaluate(fields=field, u=line)
        if x_axis == "parametric":
            ax.plot(line, vals, label="%d" % ii)

        else:
            coors = nurbs(u=line)[:, x_axis]
            ax.plot(coors, vals, label="%d" % ii)

    if legend:
        ax.legend()

    return ax

def plot_weighted_points(ax, coors, weights, min_radius=10, max_radius=50,
                         show_colorbar=False, show=False):
    """
    Plot points with given coordinates as circles/spheres with radii given by
    weights.
    """
    dim = coors.shape[1]
    ax = _get_axes(ax, dim)

    wmin, wmax = weights.min(), weights.max()
    if (wmax - wmin) < 1e-12:
        nweights = weights * max_radius / wmax

    else:
        nweights = ((weights - wmin) * (max_radius - min_radius)
                    / (wmax - wmin) + min_radius)

    if dim == 3:
        sc = ax.scatter(coors[:, 0], coors[:, 1], coors[:, 2],
                        s=nweights, c=weights, alpha=1)

    else:
        sc = ax.scatter(coors[:, 0], coors[:, 1],
                        s=nweights, c=weights, alpha=1)

    if show_colorbar:
        plt.colorbar(sc)

    if show:
        plt.show()

    return ax

def plot_iso_lines(ax, nurbs, color="b", n_points=100):
    """
    Plot the NURBS object using iso-lines in Greville abscissae coordinates.
    """
    dim = nurbs.dim
    ax = _get_axes(ax, dim)

    gas = nurbs.greville()

    if dim == 1:
        ga = gas[0]

        x0 = nm.linspace(ga[0], ga[-1], n_points)
        vals = nurbs(x0)

        if vals.shape[1] == 1:
            ax.plot(x0, vals[:, 0], color)

        else:  # Assume curve in 2D.
            ax.plot(vals[:, 0], vals[:, 1], color)

    elif dim == 2:
        ga0 = gas[0]
        ga1 = gas[1]

        x1 = nm.linspace(ga1[0], ga1[-1], n_points)
        for x0 in ga0:
            vals = nurbs(x0, x1)
            ax.plot(vals[:, 0], vals[:, 1], color)

        x0 = nm.linspace(ga0[0], ga0[-1], n_points)
        for x1 in ga0:
            vals = nurbs(x0, x1)
            ax.plot(vals[:, 0], vals[:, 1], color)

    else:
        ga0 = gas[0]
        ga1 = gas[1]
        ga2 = gas[2]

        x2 = nm.linspace(ga2[0], ga2[-1], n_points)
        for x0 in ga0:
            for x1 in ga1:
                vals = nurbs(x0, x1, x2)
                ax.plot(vals[:, 0], vals[:, 1], vals[:, 2], color)

        x1 = nm.linspace(ga1[0], ga1[-1], n_points)
        for x0 in ga0:
            for x2 in ga2:
                vals = nurbs(x0, x1, x2)
                ax.plot(vals[:, 0], vals[:, 1], vals[:, 2], color)

        x0 = nm.linspace(ga0[0], ga0[-1], n_points)
        for x1 in ga1:
            for x2 in ga2:
                vals = nurbs(x0, x1, x2)
                ax.plot(vals[:, 0], vals[:, 1], vals[:, 2], color)

    return ax

def label_points(ax, coors):
    """
    Label points with their indices.
    """
    dim = coors.shape[1]
    ax = _get_axes(ax, dim)

    shift = 0.02 * (coors.max(0) - coors.min(0))
    ccs = coors + shift
    for ic, cc in enumerate(ccs):
        ax.text(*cc, s='%d' % ic, color='b')

def plot_entities(ax, cmesh, edim, color='b', size=10):
    """
    Plot mesh topology entities using scatter plot.
    """
    coors = cmesh.get_centroids(edim)
    coors = _to2d(coors)
    dim = cmesh.dim

    ax = _get_axes(ax, dim)
    ax.scatter(*coors.T, s=size, c=color)

    return ax

def label_global_entities(ax, cmesh, edim, color='b', fontsize=10):
    """
    Label mesh topology entities using global ids.
    """
    coors = cmesh.get_centroids(edim)
    coors = _to2d(coors)
    dim = cmesh.dim

    ax = _get_axes(ax, dim)

    for ii, cc in enumerate(coors):
        ax.text(*cc.T, s=ii, color=color, fontsize=fontsize)

    return ax

def label_dofs(ax, coors, dofs, colors):
    """
    Label DOFs using the given colors.
    """
    from sfepy.postprocess.plot_dofs import _get_axes

    dim = coors.shape[1]
    ax = _get_axes(ax, dim)

    for gdof in dofs:
        cd = coors[gdof]
        ax.text(*cd.T, s='%d' % gdof,
                color=colors[gdof], fontsize=12, weight='bold')

    return ax

def plot_points(ax, points, marker, **kwargs):
    from sfepy.postprocess.plot_dofs import _get_axes

    dim = points.shape[1]
    ax = _get_axes(ax, dim)

    px, py = points[:, 0], points[:, 1]
    if dim == 2:
        ax.plot(px, py, marker, **kwargs)

    else:
        pz = points[:, 2]
        ax.plot(px, py, pz, marker, **kwargs)

    return ax

def plot_faces(ax, gel, radius, n_point, show=False):
    """
    Plot faces of a 3D geometry element as numbered oriented arcs. An arc
    centre corresponds to the first node of a face. It points from the first
    edge towards the last edge of the face.
    """
    dim = gel.dim
    ax = _get_axes(ax, dim)

    if dim < 3: return ax

    for ii, face in enumerate(gel.faces):
        cc = gel.coors[face]

        t1 = cc[1, :] - cc[0, :]
        t2 = cc[-1, :] - cc[0, :]
        n = nm.cross(t1, t2)

        nt1 = nm.linalg.norm(t1)
        nt2 = nm.linalg.norm(t2)
        angle = nm.arccos(nm.dot(t1, t2) / (nt1 * nt2))

        da = angle / (n_point - 1)

        mtx = make_axis_rotation_matrix(n, da)

        rt = cc[0] + radius * t1 / nt1
        coors = [rt]
        for ip in range(n_point - 1):
            rt = nm.dot(mtx.T, (rt - cc[0])) + cc[0]
            coors.append(rt)

        coors = nm.array(coors, dtype=nm.float64)
        centre = coors.sum(axis=0) / coors.shape[0]

        draw_arrow(ax, coors, length=0.3*radius, linewidth=3, color='r')

        if dim == 3:
            cx, cy, cz = centre
            ax.text(cx, cy, cz, ii,
                    color='r', fontsize=10, weight='light')

        else:
            cx, cy = centre
            ax.text(cx, cy, ii,
                    color='r', fontsize=10, weight='light')

    return ax