Suggestions for improving speed

edit

I'm working on visualizing a flow in R^3, and specifically focusing on how the flow behaves near a torus on inner radius 1 and outer radius 2, centered at the origin. The idea is to take a point (x,y,z), and look at this point using coordinates (r,z), in the plane containing (x,y,z) and the origin. Then I take equations in polar coordinates in this plane, and apply them to the point. Moving back into R^3, I take a rotation of the point, and add in the result of the flow restricted to the plane. Most of these conditions also require that I know the minimum distance from (x,y,z) to the torus mentioned.

This is taking a very long time, especially when I try to animate the flow. The last animation I tried took 20 hours, and that was before I found a flaw in the construction. This one works, but I'm hoping to optimize it before I try it again. Any suggestions on how to speed things up would be appreciated.

x,y,u,v = var('x y u v')

ep = 0.1 #step size
length = 350 #iterations

#initial conditions
omega_seeds = [(-1,-1,2),(1,-1,2),(-1,1,2),(1,1,2),(2,2,2),(-2,2,2),(2,-2,2),(-2,-2,2)]

twopi = 2 * RDF.pi()
c(r) = 2/3*(4**(r+1)-4)

def norm2((x,y)):
    """
    Norm of `(x,y)`
    """
    return float(sqrt((x^2 + y^2)))

def norm3((x,y,z)):
    """
    Norm squared of `(x,y,z)`
    """
    return x^2 + y^2 + z^2

def norm((x,y,z)):
    """
    Norm of `(x,y,z)`
    """
    return float(sqrt(norm3((x,y,z))))

def dist((x,y,z),(xx,yy,zz)):
    """
    Distance from `(x,y,z)` to `(xx,yy,zz)`
    """
    return norm((x - xx, y - yy,z-zz))

def atan2((x,y)): #generalized arctangent
    if x == 0 and y < 0:
        return 3*pi/2
    elif x == 0 and y > 0:
        return pi/2
    else:
        a = float(arctan(abs(y/x)))
        if x > 0:
            if y > 0:
                return a
            elif y < 0:
                return -a
            elif y == 0:
                return 0
        elif x < 0:
            if y > 0:
                return pi-a
            elif y < 0:
                return a-pi

#Functions to determine distances of a point from a specified torus
def Torus_dist((x,y,z),n):
    A = vector([0,c(n),0])
    if n%2 == 0:
        B = vector([x,y,0])
    else:
        B = vector([0,y,z])
    if A == B:
        return 4^n
    else:
        t = (2*4^n)/(norm(A - B))
        v = t*B+ (1-t)*A
        C = vector([v[0],v[1],v[2]])
    if dist(C,(x,y,z)) <= 4^n:
        return 0
    else:
        D = vector([x,y,z])
        t2 = 4^n/norm(C-D)
        v2 = t2*D + (1-t2)*C
        E = vector([v2[0],v2[1],v2[2]])
        return dist((x,y,z),E)

def dz(x,y): #vertical component of the flow restricted to the plane
    r = norm2((x,y))
    if r <=1 ...