In [1]:
cd outputs/fargo/
/home/pablo/fargo3d-0.9/outputs/fargo

FARGO3D - Advanced scripting

This notebook is an introduction to some advanced concepts for FARGO3D data reduction. We will develop some useful classes, that you could extend. You can find the code in "utils/advanced.py".

Maybe the major problem when you try to analize the FARGO3D data is the staggering of the fields. In this notebook we will see how to deal with staggered fields in an easy way.

It is important to note that FARGO3D do not dump the scalar field data with domain information. All the domain information you have is stored inside "domain_x.dat, domain_y.dat & domain_z.dat", and correspond with the edges of the cells. With it, we have all we need, but still we have to work a little bit to make it more useful.

One class we can develop is the mesh class, for storing all the mesh properties (eg: cell centers, cell edges, surfaces, volumes, etc...). In this tutorial we assume the simulation is cylindrical and 2D. Besides this, this tutorial is very general.

The surface of the cells is:

\(\displaystyle{S = \frac{\phi_{i+1}-\phi_i}{2} \left(r_{j+1}^2 - r_{j}^2\right)}\)

In [2]:
class Mesh():
    """
    Mesh class, for keeping all the mesh data.
    Input: directory [string] -> place where the domain files are.
    """
    def __init__(self, directory=""):
        if len(directory) > 1:
            if directory[-1] != '/':
                directory += '/'
        try:
            domain_x = np.loadtxt(directory+"domain_x.dat")
        except IOError:
            print "IOError with domain_x.dat"
        try:
            #We have to avoid ghost cells!
            domain_y = np.loadtxt(directory+"domain_y.dat")[3:-3]
        except IOError:
            print "IOError with domain_y.dat"
        self.xm = domain_x #X-Edge
        self.ym = domain_y #Y-Edge
        
        self.xmed = 0.5*(domain_x[:-1] + domain_x[1:]) #X-Center
        self.ymed = 0.5*(domain_y[:-1] + domain_y[1:]) #Y-Center
        
        #(Surfaces taken from the edges)
        #First we make 2D arrays for x & y, that are (theta,r)
        T,R = meshgrid(self.xm, self.ym)
        R2  = R*R
        self.surf = 0.5*(T[:-1,1:]-T[:-1,:-1])*(R2[1:,:-1]-R2[:-1,:-1])

Now, we have all the domain data easily available. Example:

In [3]:
mesh = Mesh(directory="")
print mesh.xm
[-3.14159265 -3.12932081 -3.11704896 -3.10477711 -3.09250527 -3.08023342
 -3.06796158 -3.05568973 -3.04341788 -3.03114604 -3.01887419 -3.00660234
 -2.9943305  -2.98205865 -2.96978681 -2.95751496 -2.94524311 -2.93297127
 -2.92069942 -2.90842757 -2.89615573 -2.88388388 -2.87161203 -2.85934019
 -2.84706834 -2.8347965  -2.82252465 -2.8102528  -2.79798096 -2.78570911
 -2.77343726 -2.76116542 -2.74889357 -2.73662173 -2.72434988 -2.71207803
 -2.69980619 -2.68753434 -2.67526249 -2.66299065 -2.6507188  -2.63844696
 -2.62617511 -2.61390326 -2.60163142 -2.58935957 -2.57708772 -2.56481588
 -2.55254403 -2.54027218 -2.52800034 -2.51572849 -2.50345665 -2.4911848
 -2.47891295 -2.46664111 -2.45436926 -2.44209741 -2.42982557 -2.41755372
 -2.40528188 -2.39301003 -2.38073818 -2.36846634 -2.35619449 -2.34392264
 -2.3316508  -2.31937895 -2.3071071  -2.29483526 -2.28256341 -2.27029157
 -2.25801972 -2.24574787 -2.23347603 -2.22120418 -2.20893233 -2.19666049
 -2.18438864 -2.1721168  -2.15984495 -2.1475731  -2.13530126 -2.12302941
 -2.11075756 -2.09848572 -2.08621387 -2.07394203 -2.06167018 -2.04939833
 -2.03712649 -2.02485464 -2.01258279 -2.00031095 -1.9880391  -1.97576725
 -1.96349541 -1.95122356 -1.93895172 -1.92667987 -1.91440802 -1.90213618
 -1.88986433 -1.87759248 -1.86532064 -1.85304879 -1.84077695 -1.8285051
 -1.81623325 -1.80396141 -1.79168956 -1.77941771 -1.76714587 -1.75487402
 -1.74260218 -1.73033033 -1.71805848 -1.70578664 -1.69351479 -1.68124294
 -1.6689711  -1.65669925 -1.6444274  -1.63215556 -1.61988371 -1.60761187
 -1.59534002 -1.58306817 -1.57079633 -1.55852448 -1.54625263 -1.53398079
 -1.52170894 -1.5094371  -1.49716525 -1.4848934  -1.47262156 -1.46034971
 -1.44807786 -1.43580602 -1.42353417 -1.41126232 -1.39899048 -1.38671863
 -1.37444679 -1.36217494 -1.34990309 -1.33763125 -1.3253594  -1.31308755
 -1.30081571 -1.28854386 -1.27627202 -1.26400017 -1.25172832 -1.23945648
 -1.22718463 -1.21491278 -1.20264094 -1.19036909 -1.17809725 -1.1658254
 -1.15355355 -1.14128171 -1.12900986 -1.11673801 -1.10446617 -1.09219432
 -1.07992247 -1.06765063 -1.05537878 -1.04310694 -1.03083509 -1.01856324
 -1.0062914  -0.99401955 -0.9817477  -0.96947586 -0.95720401 -0.94493217
 -0.93266032 -0.92038847 -0.90811663 -0.89584478 -0.88357293 -0.87130109
 -0.85902924 -0.84675739 -0.83448555 -0.8222137  -0.80994186 -0.79767001
 -0.78539816 -0.77312632 -0.76085447 -0.74858262 -0.73631078 -0.72403893
 -0.71176709 -0.69949524 -0.68722339 -0.67495155 -0.6626797  -0.65040785
 -0.63813601 -0.62586416 -0.61359232 -0.60132047 -0.58904862 -0.57677678
 -0.56450493 -0.55223308 -0.53996124 -0.52768939 -0.51541754 -0.5031457
 -0.49087385 -0.47860201 -0.46633016 -0.45405831 -0.44178647 -0.42951462
 -0.41724277 -0.40497093 -0.39269908 -0.38042724 -0.36815539 -0.35588354
 -0.3436117  -0.33133985 -0.319068   -0.30679616 -0.29452431 -0.28225246
 -0.26998062 -0.25770877 -0.24543693 -0.23316508 -0.22089323 -0.20862139
 -0.19634954 -0.18407769 -0.17180585 -0.159534   -0.14726216 -0.13499031
 -0.12271846 -0.11044662 -0.09817477 -0.08590292 -0.07363108 -0.06135923
 -0.04908739 -0.03681554 -0.02454369 -0.01227185  0.          0.01227185
  0.02454369  0.03681554  0.04908739  0.06135923  0.07363108  0.08590292
  0.09817477  0.11044662  0.12271846  0.13499031  0.14726216  0.159534
  0.17180585  0.18407769  0.19634954  0.20862139  0.22089323  0.23316508
  0.24543693  0.25770877  0.26998062  0.28225246  0.29452431  0.30679616
  0.319068    0.33133985  0.3436117   0.35588354  0.36815539  0.38042724
  0.39269908  0.40497093  0.41724277  0.42951462  0.44178647  0.45405831
  0.46633016  0.47860201  0.49087385  0.5031457   0.51541754  0.52768939
  0.53996124  0.55223308  0.56450493  0.57677678  0.58904862  0.60132047
  0.61359232  0.62586416  0.63813601  0.65040785  0.6626797   0.67495155
  0.68722339  0.69949524  0.71176709  0.72403893  0.73631078  0.74858262
  0.76085447  0.77312632  0.78539816  0.79767001  0.80994186  0.8222137
  0.83448555  0.84675739  0.85902924  0.87130109  0.88357293  0.89584478
  0.90811663  0.92038847  0.93266032  0.94493217  0.95720401  0.96947586
  0.9817477   0.99401955  1.0062914   1.01856324  1.03083509  1.04310694
  1.05537878  1.06765063  1.07992247  1.09219432  1.10446617  1.11673801
  1.12900986  1.14128171  1.15355355  1.1658254   1.17809725  1.19036909
  1.20264094  1.21491278  1.22718463  1.23945648  1.25172832  1.26400017
  1.27627202  1.28854386  1.30081571  1.31308755  1.3253594   1.33763125
  1.34990309  1.36217494  1.37444679  1.38671863  1.39899048  1.41126232
  1.42353417  1.43580602  1.44807786  1.46034971  1.47262156  1.4848934
  1.49716525  1.5094371   1.52170894  1.53398079  1.54625263  1.55852448
  1.57079633  1.58306817  1.59534002  1.60761187  1.61988371  1.63215556
  1.6444274   1.65669925  1.6689711   1.68124294  1.69351479  1.70578664
  1.71805848  1.73033033  1.74260218  1.75487402  1.76714587  1.77941771
  1.79168956  1.80396141  1.81623325  1.8285051   1.84077695  1.85304879
  1.86532064  1.87759248  1.88986433  1.90213618  1.91440802  1.92667987
  1.93895172  1.95122356  1.96349541  1.97576725  1.9880391   2.00031095
  2.01258279  2.02485464  2.03712649  2.04939833  2.06167018  2.07394203
  2.08621387  2.09848572  2.11075756  2.12302941  2.13530126  2.1475731
  2.15984495  2.1721168   2.18438864  2.19666049  2.20893233  2.22120418
  2.23347603  2.24574787  2.25801972  2.27029157  2.28256341  2.29483526
  2.3071071   2.31937895  2.3316508   2.34392264  2.35619449  2.36846634
  2.38073818  2.39301003  2.40528188  2.41755372  2.42982557  2.44209741
  2.45436926  2.46664111  2.47891295  2.4911848   2.50345665  2.51572849
  2.52800034  2.54027218  2.55254403  2.56481588  2.57708772  2.58935957
  2.60163142  2.61390326  2.62617511  2.63844696  2.6507188   2.66299065
  2.67526249  2.68753434  2.69980619  2.71207803  2.72434988  2.73662173
  2.74889357  2.76116542  2.77343726  2.78570911  2.79798096  2.8102528
  2.82252465  2.8347965   2.84706834  2.85934019  2.87161203  2.88388388
  2.89615573  2.90842757  2.92069942  2.93297127  2.94524311  2.95751496
  2.96978681  2.98205865  2.9943305   3.00660234  3.01887419  3.03114604
  3.04341788  3.05568973  3.06796158  3.08023342  3.09250527  3.10477711
  3.11704896  3.12932081  3.14159265]

Now we have all the mesh data, we will develop a class for working with the scalar fields taking care about the centering of the variables. Should be a good idea the use of python-inheritance for keeping the Mesh data and Parameter data inside the field. We include the Parameters class, developed in the FARGO3D tour:

In [4]:
class Parameters():
    """
    Class for reading the simulation parameters.
    input: string -> name of the parfile, normally variables.par
    """
    def __init__(self, directory=''):
        if len(directory) > 1:
            if directory[-1] != '/':
                directory += '/'
        try:
            params = open(directory+"variables.par",'r') #Opening the parfile
        except IOError:                  # Error checker.
            print  paramfile + " not found."
            return
        lines = params.readlines()     # Reading the parfile
        params.close()                 # Closing the parfile
        par = {}                       # Allocating a dictionary
        for line in lines:             #Iterating over the parfile
            name, value = line.split() #Spliting the name and the value (first blank)
            try:
                float(value)           # First trying with float
            except ValueError:         # If it is not float
                try:
                    int(value)         #                   we try with integer
                except ValueError:     # If it is not integer, we know it is string
                    value = '"' + value + '"'
            par[name] = value          # Filling the dictory
        self._params = par             # A control atribute, actually not used, good for debbuging
        for name in par:               # Iterating over the dictionary
            exec("self."+name.lower()+"="+par[name]) #Making the atributes at runtime

Now, we develop a new class, called Field class:

In [5]:
class Field(Mesh, Parameters):
    """
    Field class, it stores all the mesh, parameters and scalar data 
    for a scalar field.
    Input: field [string] -> filename of the field
           staggered='c' [string] -> staggered direction of the field. 
                                      Possible values: 'x', 'y', 'xy', 'yx'
           directory='' [string] -> where filename is
           dtype='float64' (numpy dtype) -> 'float64', 'float32', 
                                             depends if FARGO_OPT+=-DFLOAT is activated
    """
    def __init__(self, field, staggered='c', directory='', dtype='float64'):
        if len(directory) > 1:
            if directory[-1] != '/':
                directory += '/'
        Mesh.__init__(self, directory) #All the Mesh attributes inside Field!
        Parameters.__init__(self, directory) #All the Parameters attributes inside Field!

        #Now, the staggering:
        if staggered.count('x')>0:
            self.x = self.xm[:-1] #Do not dump the last element
        else:
            self.x = self.xmed
        if staggered.count('y')>0:
            self.y = self.ym[:-1]
        else:
            self.y = self.ymed
    
        self.data = self.__open_field(directory+field,dtype) #The scalar data is here.

    def __open_field(self, f, dtype):
        """
        Reading the data
        """
        field = fromfile(f, dtype=dtype)
        return field.reshape(self.ny, self.nx)
    
    """                                                                                                  
    Now we can include a propper function for making fancy plots!  We                                    
    will use the **karg arguments + matplotlib. There are two ways to                                    
    plot a field. The best option is "imshow" (an amazing quick                                          
    function), but with "imshow" is hard to handle cartesian                                             
    coordinates from a cylindrical simulation (actually it is not so                                     
    hard, you could use the gridddata function!).  The best choise for                                   
    doing quick cartesian plots from a cylindrical data is the use of                                    
    pcolormesh, but it is a tortoise; if your data is too big you                                        
    should consider the use gridddata.                                                                    
    """

    def plot(self, log=False, cartesian=False, cmap='cubehelix', **karg):
        """
        A layer for plt.imshow or pcolormesh function.
        if cartesian = True, pcolormesh is launched.
        """
        ax = gca()
        if log:
            data = np.log(self.data)
        else:
            data = self.data
        if cartesian:
            T,R = meshgrid(self.x,self.y)
            X = R*cos(T)
            Y = R*sin(T)
            pcolormesh(X,Y,data,cmap=cmap,**karg)
        else:
            ax.imshow(data, cmap = cmap, origin='lower',aspect='auto',
                      extent=[self.x[0],self.x[-1],self.y[0],self.y[-1]],
                      **karg)
            
    #The same but for contours
    def contour(self, log=False, cartesian=False, **karg):
        if log:
            data = np.log(self.data)
        else:
            data = self.data
        ax = gca()
        T,R = meshgrid(self.x,self.y)
        if cartesian:
            X = R*cos(T)
            Y = R*sin(T)
            contour(X,Y,data,**karg)
        else:
            contour(T,R,data,**karg)

We are developed very simple classes that simplify the work with the FARGO3D outputs. Now, it is time for tests! For the tests below, we are in "outpus/fargo" directory:

In [6]:
rho = Field("gasdens10.dat")
In [7]:
figure(figsize=(7,7))
rho.contour(cartesian=True,log=True,
            levels=linspace(log(rho.data).min(),log(rho.data).max(),20),
            colors='k', alpha=0.3, linestyles='solid')
rho.plot(cartesian=True, log=True, vmax=-5.0)
ax = gca()
ax.xaxis.set_visible(False)
ax.yaxis.set_visible(False)
t = linspace(-pi,pi,1000)
plot(cos(t),sin(t),'w')
xlim(0.7,1.3)
ylim(-.35,0.35)
Out[7]:
(-0.35, 0.35)

Now the same but in cylindrical cordinates!

In [8]:
figure(figsize=(7,7))
rho.plot(log=True,interpolation='nearest')
rho.contour(log=True, levels=linspace(log(rho.data).min(),log(rho.data).max(),30),
            colors='k', alpha=0.3, linestyles='solid')
ax = gca()
ax.xaxis.set_visible(False)
ax.yaxis.set_visible(False)
tmax = arctan2(0.35,1.3)
ylim(0.7,1.3)
xlim(-tmax,tmax)
Out[8]:
(-0.26299473168091947, 0.26299473168091947)

All right! Now it is time to test what happens if the fields are staggered:

In [9]:
vx = Field("gasvx10.dat",staggered='x')
vy = Field("gasvy10.dat",staggered='y')

might be desirable to have some method for centering the staggered fields in some direction (it is useful when you want to derive some new field, for example, the total velocity). We will develop a funcion that takes a Field class and make the centering. You could include this function inside the Field class.

In [10]:
def shift_field(Field,direction):
    import copy
    """
    Half cell shifting along the direction provided by direction 
    direction can be ('x','y', 'xy', 'yx').
    
    After a call of this function, Field.xm/xmed has not
    sense anymore (it is not hard to improve).
    """
    F = copy.deepcopy(Field)
    if direction.count('x')>0:
        F.data = 0.5*(Field.data[:,1:]+Field.data[:,:-1])
        F.x = 0.5*(Field.x[1:]+Field.x[:-1])
    if direction.count('y')>0:
        F.data = 0.5*(F.data[1:,:]+F.data[:-1,:])
        F.y = 0.5*(F.y[1:]+F.y[:-1])

    F.nx = len(F.x)
    F.ny = len(F.y)
    
    return F
In [11]:
vxc = shift_field(vx,'x')

We can see that the field was properly centered in X comparing vx.x and rho.x:

In [12]:
any(vxc.x-rho.x[:-1])
Out[12]:
False

Ok, now the centering of vy:

In [13]:
vyc = shift_field(vy,'y')
In [14]:
any(vyc.y-rho.y[:-1])
Out[14]:
False

Unfortunately, the dimensions of vy, vx and rho are not the same. We need a function for cutting fields (again, you could include it inside the Field class):

In [15]:
def cut_field(Field, direction, side):
    """                                                                                                  
    Cutting a field:
    Input: field --> a Field class                                                                       
           axis  --> 'x', 'y' or 'xy'                                                                    
           side  --> 'p' (plus), 'm' (minnus), 'pm' (plus/minnus)                                        
    """                                                                                                  
    import copy
    
    cutted_field = copy.deepcopy(Field)                                                                     
    ny,nx = Field.ny, Field.nx
    mx = my = px = py = 0
    
    if direction.count('x')>0:
        if side.count('m')>0:
            mx = 1                                                                                       
        if side.count('p')>0:
            px = 1
    if direction.count('y')>0:
        if side.count('m')>0:
            my = 1
        if side.count('p')>0:
            py = 1
                                                                                                         
    cutted_field.data = Field.data[my:ny-py,mx:nx-px]                                                       
    cutted_field.x = cutted_field.x[mx:nx-px]                                                                  
    cutted_field.y = cutted_field.y[my:ny-py]                                                                  
                                                                                                         
    return cutted_field
In [16]:
vxcc = cut_field(vxc,direction='y',side='p')
vycc = cut_field(vyc,direction='x',side='p')
rhoc = cut_field(rho,direction='xy', side='p')
In [17]:
print vxcc.data.shape
print vycc.data.shape
print rhoc.data.shape
(255, 511)
(255, 511)
(255, 511)

As you can see, now de dimesion of each field is the same. All right, now it is time to compute the total velocity scalar field:

In [18]:
import copy
v = copy.deepcopy(vxcc) #A very useful trick!
v.data = sqrt(vxcc.data**2+vycc.data**2)
In [19]:
figure(figsize=(9,7))
v.plot(log=True, cartesian=True, vmin=-4.5); colorbar()
ax = gca()
ax.xaxis.set_visible(False)
ax.yaxis.set_visible(False)
title("Velocity", size=25)
xlim(0.8,1.2)
ylim(-.2,0.2)
Out[19]:
(-0.2, 0.2)

Ok, you could say, is the same as in the FARGO3D tour! It is true, but now you have a set of tools for doing that a lot of times with less chances of a mistake.

If you want a quick plot of the velocity field, you could do some similar to:

In [20]:
def vector_field(vx,vy, **karg):
    nsx = nsy = 3
    T,R = meshgrid(vx.x[::nsx],vx.y[::nsy])
    X = R*cos(T)
    Y = R*sin(T)
    vx = vx.data[::nsy,::nsx]
    vy = vy.data[::nsy,::nsx]
    U = vy*cos(T) - vx*sin(T)
    V = vy*sin(T) + vx*cos(T)
    ax = gca()
    ax.quiver(X,Y,U,V,scale=5,pivot='midle', **karg)
In [21]:
figure(figsize=(9,7))
rho.plot(log=True, cartesian=True); colorbar()
vector_field(vxcc, vycc)
ax = gca()
ax.xaxis.set_visible(False)
ax.yaxis.set_visible(False)
title("Velocity", size=25)
xlim(0.8,1.2)
ylim(-.2,0.2)
Out[21]:
(-0.2, 0.2)

Useful example: Computing the vortencity.

As a real-life example, we will show you how to compute the vortencity (\(\omega/\rho = \nabla\times\vec{v}|_z/\rho\)) in a standard simulation. In cylindrical coordinates, the vortencity is defined as: \((\omega/\rho)\)=\(\displaystyle{\frac{1}{\rho}\partial_r(rV_\phi) - \partial_\phi(V_r)}\), (\(V\phi\) is in the inertial frame, be careful with OMEGAFRAME parameter!).

In [22]:
def vortencity(rho, vx, vy):
    """
    Computing the vortencity from a staggered velocity field.
    """
    Tx,Rx = np.meshgrid(vx.x,vx.y)
    Ty,Ry = np.meshgrid(vy.x,vy.y)

    rvx = Rx*(vx.data)

    curldata = ((rvx[1:,1:]-rvx[:-1,1:])/(Rx[1:,1:]-Rx[:-1,1:]) -
                (vy.data[1:,1:]-vy.data[1:,:-1])/(Ty[1:,1:]-Ty[1:,:-1]))

    curl = copy.deepcopy(vx)
    curl.nx = curl.nx-1
    curl.ny = curl.ny-1
    curl.x    = vx.x[:-1]
    curl.y    = vy.y[:-1]

    rho_corner = 0.25*(rho.data[1:,1:]+rho.data[:-1,:-1]+
                       rho.data[1:,:-1]+rho.data[:-1,1:])

    T,R = np.meshgrid(curl.x,curl.y)
    curl.data = curldata/R
    
    curl.data = (curl.data + 2.0*rho.omegaframe)/rho_corner
    
    return curl
In [23]:
vort = vortencity(rho, vx, vy)
In [24]:
figure(figsize=(7,7))
vort.plot(cartesian=True)
title("Vortencity", size=25)
xlim(0.6,1.4)
ylim(-.4,0.4)
Out[24]:
(-0.4, 0.4)
In [25]:
figure(figsize=(7,7))
vort.plot()
title("Vortencity", size=25)
ylim(0.7,1.3)
Out[25]:
(0.7, 1.3)

The last part of this notebook is for computing some streamlines! We will put here some very simple functions (an integrator, an interpolator and a recursive stream mannager). These funtions are very simple and we will focus on its utilization:

In [26]:
def euler(vx, vy, x, y, reverse):
    """
    Euler integrator for computing the streamlines.
    Parameters:
    ----------

    x,y: Float.
         starter coordinates.
    reverse: Boolean.
             If reverse is true, the integratin step is negative.
    Reverse inverts the sign of the velocity

    Output
    ------

    (dx,dy): (float,float).
             Are the azimutal and radial increment.
             Only works for cylindrical coordinates.
    """
    sign = 1.0
    if reverse:
        sign = -1
    vphi = get_v(vx,x,y)
    vrad = get_v(vy,x,y)
    if vphi == None or vrad == None: #Avoiding problems...                                                                                                                                                                                    
        return None,None
    l = np.min((((vx.xmax-vx.xmin)/vx.nx),((vx.ymax-vx.ymin)/vx.ny)))
    h = 0.5*l/np.sqrt((vphi**2+vrad**2))

    return sign*h*np.array([vphi/y,vrad])


def bilinear(x,y,f,p):
    """
    Computing bilinear interpolation.
    Parameters
    ----------
    x = (x1,x2); y = (y1,y2)
    f = (f11,f12,f21,f22)
    p = (x,y)
    where x,y is the interpolated point and
    fij is the value of the function at the
    point (xi,yj).
    Output
    ------
    f(p): Float.
          The interpolated value of the function f(p) = f(x,y)
    """

    xp  = p[0]; yp   = p[1]; x1  = x[0]; x2  = x[1]
    y1  = y[0]; y2  = y[1];  f11 = f[0]; f12 = f[1]
    f21 = f[2]; f22 = f[3]
    t = (xp-x1)/(x2-x1);    u = (yp-y1)/(y2-y1)

    return (1.0-t)*(1.0-u)*f11 + t*(1.0-u)*f12 + t*u*f22 + u*(1-t)*f21

def get_v(v, x, y):
    """
    For a real set of coordinates (x,y), returns the bilinear            
    interpolated value of a Field class.
    """

    i = int((x-v.xmin)/(v.xmax-v.xmin)*v.nx)
    j = int((y-v.ymin)/(v.ymax-v.ymin)*v.ny)

    if i<0 or j<0 or i>v.data.shape[1]-2 or j>v.data.shape[0]-2:
        return None

    f11 = v.data[j,i]
    f12 = v.data[j,i+1]
    f21 = v.data[j+1,i]
    f22 = v.data[j+1,i+1]
    try:
        x1  = v.x[i]
        x2  = v.x[i+1]
        y1  = v.y[j]
        y2  = v.y[j+1]
        return bilinear((x1,x2),(y1,y2),(f11,f12,f21,f22),(x,y))
    except IndexError:
        return None
    
    
    
def get_stream(vx, vy, x0, y0, nmax=1000000, maxlength=4*np.pi, bidirectional=True, reverse=False):
    """
    Function for computing a streamline.
    Parameters:
    -----------

    x0,y0: Floats.
          Initial position for the stream
    nmax: Integer.
          Maxium number of iterations for the stream.
    maxlength: Float
               Maxium allowed length for a stream
    bidirectional=True
                  If it's True, the stream will be forward and backward computed.
    reverse=False
            The sign of the stream. You can change it mannualy for a single stream,
            but in practice, it's recommeneded to use this function without set reverse
            and setting bidirectional = True.

    Output:
    -------

    If bidirectional is False, the function returns a single array, containing the streamline:
    The format is:

                                      np.array([[x],[y]])

    If bidirectional is True, the function returns a tuple of two arrays, each one with the same
    format as bidirectional=False.
    The format in this case is:

                            (np.array([[x],[y]]),np.array([[x],[y]]))

    This format is a little bit more complicated, and the best way to manipulate it is with iterators.
    For example, if you want to plot the streams computed with bidirectional=True, you should write
    some similar to:

    stream = get_stream(x0,y0)
    ax.plot(stream[0][0],stream[0][1]) #Forward
    ax.plot(stream[1][0],stream[1][1]) #Backward

    """
    
    if bidirectional:
        s0 = get_stream(vx, vy, x0, y0, reverse=False, bidirectional=False, nmax=nmax,maxlength=maxlength)
        s1 = get_stream(vx, vy, x0, y0, reverse=True,  bidirectional=False, nmax=nmax,maxlength=maxlength)
        return (s0,s1)

    l = 0
    x = [x0]
    y = [y0]

    for i in xrange(nmax):
        ds = euler(vx, vy, x0, y0, reverse=reverse)
        if ds[0] == None:
            if(len(x)==1):
                print "There was an error getting the stream, ds was NULL (see get_stream)."
            break
        l += np.sqrt(ds[0]**2+ds[1]**2)
        dx = ds[0]
        dy = ds[1]
        if(np.sqrt(dx**2+dy**2)<1e-13):
            print "Warning (get_stream): ds is very small, maybe you're in a stagnation point."
            print "Try selecting another initial point."
            break
        if l > maxlength:
            print "maxlength reached: ", l
            break
        x0 += dx
        y0 += dy
        x.append(x0)
        y.append(y0)
    return np.array([x,y])

def get_random_streams(vx, vy, xmin=None, xmax=None, ymin=None, ymax=None, n=30, nmax=100000):
    if xmin == None:
        xmin = vx.xmin
    if ymin == None:
        ymin = vx.ymin
    if xmax == None:
        xmax = vx.xmax
    if ymax == None:
        ymax = vx.ymax
    X = xmin + np.random.rand(n)*(xmax-xmin)
    Y = ymin + np.random.rand(n)*(ymax-ymin)
    streams = []
    counter = 0
    for x,y in zip(X,Y):
        stream = get_stream(vx, vy, x, y, nmax=nmax, bidirectional=True)
        streams.append(stream)
        counter += 1
        print "stream ",counter, "done"
    return streams

def plot_random_streams(streams, cartesian=False, **kargs):
    ax = plt.gca()
    print np.shape(streams)
    for stream in streams:
        for sub_stream in stream:
            if cartesian:
                ax.plot(sub_stream[1]*cos(sub_stream[0]),sub_stream[1]*sin(sub_stream[0]),**kargs)
            else:
                ax.plot(sub_stream[0],sub_stream[1],**kargs)
In [27]:
random_streams = get_random_streams(vx,vy, n=50)
stream  1 done
stream  2 done
stream  3 done
stream  4 done
stream  5 done
stream  6 done
stream  7 done
stream  8 done
stream  9 done
There was an error getting the stream, ds was NULL (see get_stream).
There was an error getting the stream, ds was NULL (see get_stream).
stream  10 done
stream  11 done
stream  12 done
stream  13 done
stream  14 done
stream  15 done
stream  16 done
stream  17 done
stream  18 done
stream  19 done
stream  20 done
stream  21 done
stream  22 done
stream  23 done
stream  24 done
stream  25 done
stream  26 done
stream  27 done
stream  28 done
stream  29 done
stream  30 done
stream  31 done
stream  32 done
stream  33 done
stream  34 done
stream  35 done
stream  36 done
stream  37 done
stream  38 done
stream  39 done
stream  40 done
stream  41 done
stream  42 done
stream  43 done
stream  44 done
stream  45 done
stream  46 done
stream  47 done
stream  48 done
stream  49 done
stream  50 done

In [28]:
figure(figsize=(9,9))
rho.plot(log=True, cartesian=True)
vector_field(vx, vy, color='w', alpha=0.5)
plot_random_streams(random_streams,cartesian=True,color='k', linewidth=2, alpha=0.5)
ax = gca()
ax.xaxis.set_visible(False)
ax.yaxis.set_visible(False)
title("Velocity", size=25)
xlim(0.5,1.4)
ylim(-0.45,0.45)
(50, 2, 2)

Out[28]:
(-0.45, 0.45)