//  [python]
def streamfunction(om, r, psi, ur, uz, dr, dz, 
        BC_Walls={'south':BC_Dir_0, 'west':BC_Dir_Uz0, 'north':BC_Neu_uz, 'east':BC_Neu_ur0},
        nIter=300000, rTol=1e-12, verbose=False):
    """ 
    
    """
    psi2 = np.copy(psi)
    nz,nr= psi.shape
    
    idz2 = 1/(dz*dz)
    idr2 = 1/(dr*dr)
    for it in range(nIter):   
        psi2[1:nz-1,1:nr-1] = (om[1:nz-1,1:nr-1]*r[1:nr-1]  
                             + idz2*(psi[0:nz-2,1:nr-1]+psi[2:nz,1:nr-1])
                             + idr2*(psi[1:nz-1,0:nr-2]+psi[1:nz-1,2:nr])
                             - 1/(2*dr*r[1:nr-1])*(psi[1:nz-1,2:nr]-psi[1:nz-1,0:nr-2])
                             )/(2*(idz2+idr2))        
                            
        #replace values on boundary nodes with previous values
        #psi2[node_type>0] = psi[node_type>0]
        # --- Inlet (i=0, j)
        if BC_Walls['west']==BC_Dir_Uz0:
            # - Dirichlet - uz known and constant
            psi2[0,:] = 0.5*uz[0,0]*r**2
        elif BC_Walls['west']==BC_Dir_uz:
            # - Dirichlet - uz known
            psi2[0,1:] = sciint.cumtrapz(r*uz[0,:], r)
        elif BC_Walls['west']==BC_Neu_ur0:
            # - Neumann - No radial flow: ur=0,  dpsi/dz = -r ur = 0 - First order
            #psi2[0,:] = psi2[1,:] 
            # - Neumann - No radial flow: ur=0,  dpsi/dz = -r ur = 0 - Second order
            psi2[0,:] = 1/3*(-psi2[2,:]+4*psi2[1,:])
        elif BC_Walls['west']==BC_Neu_ur:
            # - Neumann - ur known,  dpsi/dz = -ur r  - First order
            #psi2[0,:] = psi2[1,:] + dz*ur[0,:]*r[:]
            # - Neumann - ur known,  dpsi/dz = -ur r  - Second order forward
            psi2[0,:] = 1/3*(-psi2[2,:]+4*psi2[1,:] + 2*dz*ur[0,:]*r[:])
        else:
            raise NotImplementedError('BC_walls west', BC_Walls)
        

        # --- Outlet (i=nz-1, j)
        if BC_Walls['east']==BC_Dir_Uz0:
            # - Dirichlet - uz known and constant
            psi2[nz-1,:] = 0.5*uz[0,0]*r**2
        elif BC_Walls['east']==BC_Dir_uz:
            # - Dirichlet - uz known
            psi2[n1-1,1:] = sciint.cumtrapz(r*uz[nz-1,:], r)
        elif BC_Walls['east']==BC_Neu_ur0:
            # - Neumann - no radial flow ur=0 - First order
            #psi2[nz-1,:] = psi2[nz-2,:]
            # - Neumann - no radial flow ur=0 - Second order
            psi2[nz-1,:] = 1/3*(-psi2[nz-3,:] + 4*psi2[nz-2,:])
        elif BC_Walls['east']==BC_Neu_ur:
            # - Neumann - ur known:  dpsi/dz = -ur r  - First order
            #psi2[nz-1,:] = psi2[nz-2,:] - dz*ur[nz-1,:]*r[:]
            # - Neumann - ur known:  dpsi/dz = -ur r  - Second order
            psi2[nz-1,:] = 1/3*(-psi2[nz-3,:] + 4*psi2[nz-2,:] - 2*dz*ur[nz-1,:]*r[:])
        #for j in range(nr):
        #    if (node_type[nz-1,j]!=WALL):
        #        psi2[nz-1,j] = psi2[nz-2,j] - dz*(v[nz-1,j])*r[j]
        else:
            raise NotImplementedError('BC_walls east', BC_Walls)

        # --- Top boundary (i, j=nr-1)
        if BC_Walls['north']==BC_Dir_Uz0:
            # - Dirichlet - freestream constant velocity u0
            psi2[:,nr-1] = 0.5*uz[0,0]*r[nr-1]**2
        elif BC_Walls['north']==BC_Dir_uz:
            # - Dirichlet - freestream varying velocity
            psi2[:,nr-1] = np.trapz(r*uz[0,:],r)
        elif BC_Walls['north']==BC_Neu_uz:
            # - Neumann - uz known:  dpsi/dr = r uz - First order
            #psi2[:,nr-1] = psi2[:,nr-2] + dr*uz[:,nr-1]*r[nr-1]
            # - Neumann - uz known:  dpsi/dr = r uz - Second order
            psi2[:,nr-1] = 1/3*(-psi2[:,nr-3] + 4*psi2[:,nr-2] + 2*dr*uz[:,nr-1]*r[nr-1])
        elif BC_Walls['north']==BC_Neu_uz0:
            # - Neumann - dpsi/dr = ruz=0 - First order
            #psi2[:,nr-1] = psi2[:,nr-2]
            # - Neumann - dpsi/dr = ruz=0 - Second order
            psi2[:,nr-1] = 1/3*(-psi2[:,nr-3] + 4*psi2[:,nr-2])
        else:
            raise NotImplementedError('BC_walls south', BC_Walls)

        # --- Bottom boundary (i, j=0)
        if BC_Walls['south']==BC_Dir_0:
            # - Dirichlet - psi known to be 0
            psi2[:,0] = 0
        else:
            raise NotImplementedError('BC_walls', BC_Walls)
 
        # copy back solution        
        psi = np.copy(psi2) 
            
        # --- Convergence check
        if (it%25==0):
            R = np.zeros_like(psi)
            R[1:nz-1,1:nr-1] = (om[1:nz-1,1:nr-1]*r[1:nr-1] + 
                    idz2*(psi[0:nz-2,1:nr-1]-2*psi[1:nz-1,1:nr-1]+psi[2:nz,1:nr-1])+
                    idr2*(psi[1:nz-1,0:nr-2]-2*psi[1:nz-1,1:nr-1]+psi[1:nz-1,2:nr])-
                    1/(2*dr*r[1:nr-1])*(psi[1:nz-1,2:nr]-psi[1:nz-1,0:nr-2]))
            #R[node_type>0] = 0
            norm     = np.linalg.norm(R)
            norm_avg = norm/(nz*nr)
            if (norm_avg<rTol): 
                if verbose:
                    print('>>> Converged, number of iterations',it)
                return psi2
            if (it%1000==0) and verbose:
                print('    it: {:07d} norm: {:6.2e}   {:6.2e}'.format(it, norm, norm_avg))
    print('[FAIL] Psi failed to converge, norm = {:6.2e}'.format(norm))
    return psi