Simple routines that imitate MATLAB-style "easy plotting" commands.

np_plots.py

import numpy as np
import math as m

# calculates point of parameterized curve
# accepts parameterization interval, parameterization function
# and discretization step
def ezLinePlot(Boundaries, CurveParametrization, Steps):
    left, right = Boundaries
    curveParamPoints = np.linspace(left, right, Steps)
    curvePoints = [CurveParametrization(t) for t in curveParamPoints]
    X, Y = zip(*curvePoints)
    return [X, Y]
#
# calculates vector field on curve
# accepts parameterization interval, curve equation,
# vector field equation and parameterization step
def ezLineQuiver(Boundaries, CurveParametrization, VectorField, Steps, Normalize=False):
    curvePoints = ezLinePlot(Boundaries, CurveParametrization, Steps) # returns pair of lists
    X, Y = curvePoints
    vectorField = [VectorField([X[i],Y[i]]) for i in xrange(len(X))]
    if Normalize:
        vectorField = [np.array(v)/m.sqrt(np.dot(v,v)) for v in vectorField]
    U, V = zip(*vectorField)
    return [X, Y, U, V]

# calculates vector field in 3D and 2D domains
# accepts array of boundaries, array of discretisation steps,
# vector field function
def ezDomainQuiver3D (Boundaries, Steps, VectorField, Normalize=False):
    assert(len(Steps)==len(Boundaries))
    assert(len(Boundaries[0])==2)
    xgrid, ygrid, zgrid = [np.linspace(Boundaries[i][0], Boundaries[i][1], Steps[i]) for i in xrange(len(Steps))]
    points = [[x, y, z] for x in xgrid for y in ygrid for z in zgrid]
    vectors = [VectorField(p) for p in points]
    X, Y, Z = zip(*points)
    U, V, W = zip(*vectors)
    return [X, Y, Z, U, V, W]

def ezDomainQuiver2D (Boundaries, Steps, VectorField, Normalize=False):
    assert(len(Steps)==len(Boundaries))
    assert(len(Boundaries[0])==2)
    xgrid, ygrid = [np.linspace(Boundaries[i][0], Boundaries[i][1], Steps[i]) for i in xrange(len(Steps))]
    points = [[x, y] for x in xgrid for y in ygrid]
    vectors = [VectorField(p) for p in points]
    if Normalize:
        vectors = [np.array(v)/m.sqrt(np.dot(v,v)) for v in vectors]
    X, Y = zip(*points)
    U, V = zip(*vectors)
    return [X, Y, U, V]