import math
import numpy
from matplotlib import pyplot


def velocity_grid(strength, x_source, y_source, X, Y):
    """
    Parameters
    ----------
    strength: list[floats] or float
        Strength of the sources/sinks
    x_source: list[floats] or float
        x-coordinates of the sources and sinks.
    y_source: list[floats] or float
        y-coordinates of the sources and sinks.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.

    Returns
    -------
    u_source 2D Numpy array of floats
        x-velocity field
    v_source 2D Numpy array of floats
        y-velocity field

    Note
    ----
    sources are superimposed into one velocity field
    """
    if isinstance(x_source, list):
        u_source = [0.0] * len(x_source)
        v_source = [0.0] * len(y_source)
        for i in range(0, len(x_source)):
            u_source[i] = strength[i] / (2 * math.pi) * (X - x_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
            v_source[i] = strength[i] / (2 * math.pi) * (Y - y_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
        u_source = sum(u_source)
        v_source = sum(v_source)
    else:
        u_source = strength / (2 * math.pi) * (X - x_source) / ((X - x_source) ** 2 + (Y - y_source) ** 2)
        v_source = strength / (2 * math.pi) * (Y - y_source) / ((X - x_source) ** 2 + (Y - y_source) ** 2)
    return u_source, v_source


def get_stream_function(strength, x_source, y_source, X, Y):
    """
    Returns the stream-function generated by a source/sink.
    
    Parameters
    ----------
    strength: float
        Strength of the source/sink.
    xs: float
        x-coordinate of the source (or sink).
    ys: float
        y-coordinate of the source (or sink).
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    if isinstance(x_source, list):
        psi = [0.0] * len(x_source)
        for i in range(0, len(x_source)):
            psi[i] = strength[i] / (2 * numpy.pi) * numpy.arctan2((Y - y_source[i]), (X - x_source[i]))
        psi = sum(psi)
    else:
        psi = strength / (2 * numpy.pi) * numpy.arctan2((Y - y_source), (X - x_source))

    return psi

def get_velocity_doublet(strength, xd, yd, X, Y):
    """
    Returns the velocity field generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    u: 2D Numpy array of floats
        x-component of the velocity vector field.
    v: 2D Numpy array of floats
        y-component of the velocity vector field.
    """
    if isinstance(xd, list):
        u = [0.0] * len(xd)
        v = [0.0] * len(yd)
        for i in range(0, len(xd)):
            u[i] = - strength[i] / (2 * math.pi) * ((X - xd[i]) ** 2 - (Y - yd[i]) ** 2) / ((X - xd[i]) ** 2 + (
                Y - yd[i]) ** 2) ** 2
            v[i] = - strength[i] / (2 * math.pi) * 2 * (X - xd[i]) * (Y - yd[i]) / ((X - xd[i]) ** 2 + (
                Y - yd[i]) ** 2) ** 2
        u = sum(u)
        v = sum(v)
    else:
        u = - strength / (2 * math.pi) * ((X - xd) ** 2 - (Y - yd) ** 2) / ((X - xd) ** 2 + (Y - yd) ** 2) ** 2
        v = - strength / (2 * math.pi) * 2 * (X - xd) * (Y - yd) / ((X - xd) ** 2 + (Y - yd) ** 2) ** 2

    return u, v

def get_stream_function_doublet(strength, xd, yd, X, Y):
    """
    Returns the stream-function generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    psi = - strength / (2 * math.pi) * (Y - yd) / ((X - xd) ** 2 + (Y - yd) ** 2)

    return psi

# establish grid
N = 50
x_start, x_end = 0.0, 3.0
y_start, y_end = 0.0, 2.0
x = numpy.linspace(x_start, x_end, N)
y = numpy.linspace(y_start, y_end, N)
# mesh creation
X, Y = numpy.meshgrid(x, y)

# source + sink input information
"""
N_source = 100
strength = [0.1] * N_source # source strength
x_source = list(numpy.linspace(-0.5, 1, N_source)) #x_coordinates
y_source = list(numpy.linspace(0.5, 0.0, N_source)) 
#y_source = [0.0] * N_source # y_coordinates
"""
strength = [5.0, -1.0]
x_source = [1.0, 1.5]
y_source = [1.0, 1.0]

# adding in free stream velocity field
u_inf = 2.0

u_freestream = u_inf * numpy.ones((N, N), dtype=float)
v_freestream = numpy.zeros((N, N), dtype=float)

# compute the stream function
psi_freestream = u_inf * Y

# compute velocity grid
u_source, v_source = velocity_grid(strength, x_source, y_source, X, Y)
psi_source = get_stream_function(strength, x_source, y_source, X, Y)

# compute final velocity grid
u = u_freestream + u_source
v = v_freestream + v_source
psi = psi_freestream + psi_source


# plot the gird of points
size = 10
pyplot.figure(figsize=(size, (y_end-y_start)/(x_end-x_start)*size))
pyplot.xlabel('x', fontsize = 16)
pyplot.ylabel('y', fontsize = 16)
pyplot.xlim(x_start, x_end)
pyplot.ylim(y_start, y_end)
pyplot.streamplot(X, Y, u, v, density=2, linewidth=1, arrowsize=2, arrowstyle='->')
pyplot.scatter(x_source, y_source, s=10, color = '#CD2305', marker = 'o', linewidth = 0)
pyplot.contour(X,Y,psi,
               #levels=[-strength[0]/2.,0, +strength[0]/2.],
               levels=[1,3],
               colors='#CD2305',linewidths=2, linestyles='solid')
pyplot.show()

54fa0fd06a221659322ff47ef01fad0c24e05dee
<mask token>


def velocity_grid(strength, x_source, y_source, X, Y):
    """
    Parameters
    ----------
    strength: list[floats] or float
        Strength of the sources/sinks
    x_source: list[floats] or float
        x-coordinates of the sources and sinks.
    y_source: list[floats] or float
        y-coordinates of the sources and sinks.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.

    Returns
    -------
    u_source 2D Numpy array of floats
        x-velocity field
    v_source 2D Numpy array of floats
        y-velocity field

    Note
    ----
    sources are superimposed into one velocity field
    """
    if isinstance(x_source, list):
        u_source = [0.0] * len(x_source)
        v_source = [0.0] * len(y_source)
        for i in range(0, len(x_source)):
            u_source[i] = strength[i] / (2 * math.pi) * (X - x_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
            v_source[i] = strength[i] / (2 * math.pi) * (Y - y_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
        u_source = sum(u_source)
        v_source = sum(v_source)
    else:
        u_source = strength / (2 * math.pi) * (X - x_source) / ((X -
            x_source) ** 2 + (Y - y_source) ** 2)
        v_source = strength / (2 * math.pi) * (Y - y_source) / ((X -
            x_source) ** 2 + (Y - y_source) ** 2)
    return u_source, v_source


def get_stream_function(strength, x_source, y_source, X, Y):
    """
    Returns the stream-function generated by a source/sink.
    
    Parameters
    ----------
    strength: float
        Strength of the source/sink.
    xs: float
        x-coordinate of the source (or sink).
    ys: float
        y-coordinate of the source (or sink).
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    if isinstance(x_source, list):
        psi = [0.0] * len(x_source)
        for i in range(0, len(x_source)):
            psi[i] = strength[i] / (2 * numpy.pi) * numpy.arctan2(Y -
                y_source[i], X - x_source[i])
        psi = sum(psi)
    else:
        psi = strength / (2 * numpy.pi) * numpy.arctan2(Y - y_source, X -
            x_source)
    return psi


<mask token>


def get_stream_function_doublet(strength, xd, yd, X, Y):
    """
    Returns the stream-function generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    psi = -strength / (2 * math.pi) * (Y - yd) / ((X - xd) ** 2 + (Y - yd) ** 2
        )
    return psi


<mask token>

<mask token>


def velocity_grid(strength, x_source, y_source, X, Y):
    """
    Parameters
    ----------
    strength: list[floats] or float
        Strength of the sources/sinks
    x_source: list[floats] or float
        x-coordinates of the sources and sinks.
    y_source: list[floats] or float
        y-coordinates of the sources and sinks.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.

    Returns
    -------
    u_source 2D Numpy array of floats
        x-velocity field
    v_source 2D Numpy array of floats
        y-velocity field

    Note
    ----
    sources are superimposed into one velocity field
    """
    if isinstance(x_source, list):
        u_source = [0.0] * len(x_source)
        v_source = [0.0] * len(y_source)
        for i in range(0, len(x_source)):
            u_source[i] = strength[i] / (2 * math.pi) * (X - x_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
            v_source[i] = strength[i] / (2 * math.pi) * (Y - y_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
        u_source = sum(u_source)
        v_source = sum(v_source)
    else:
        u_source = strength / (2 * math.pi) * (X - x_source) / ((X -
            x_source) ** 2 + (Y - y_source) ** 2)
        v_source = strength / (2 * math.pi) * (Y - y_source) / ((X -
            x_source) ** 2 + (Y - y_source) ** 2)
    return u_source, v_source


def get_stream_function(strength, x_source, y_source, X, Y):
    """
    Returns the stream-function generated by a source/sink.
    
    Parameters
    ----------
    strength: float
        Strength of the source/sink.
    xs: float
        x-coordinate of the source (or sink).
    ys: float
        y-coordinate of the source (or sink).
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    if isinstance(x_source, list):
        psi = [0.0] * len(x_source)
        for i in range(0, len(x_source)):
            psi[i] = strength[i] / (2 * numpy.pi) * numpy.arctan2(Y -
                y_source[i], X - x_source[i])
        psi = sum(psi)
    else:
        psi = strength / (2 * numpy.pi) * numpy.arctan2(Y - y_source, X -
            x_source)
    return psi


def get_velocity_doublet(strength, xd, yd, X, Y):
    """
    Returns the velocity field generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    u: 2D Numpy array of floats
        x-component of the velocity vector field.
    v: 2D Numpy array of floats
        y-component of the velocity vector field.
    """
    if isinstance(xd, list):
        u = [0.0] * len(xd)
        v = [0.0] * len(yd)
        for i in range(0, len(xd)):
            u[i] = -strength[i] / (2 * math.pi) * ((X - xd[i]) ** 2 - (Y -
                yd[i]) ** 2) / ((X - xd[i]) ** 2 + (Y - yd[i]) ** 2) ** 2
            v[i] = -strength[i] / (2 * math.pi) * 2 * (X - xd[i]) * (Y - yd[i]
                ) / ((X - xd[i]) ** 2 + (Y - yd[i]) ** 2) ** 2
        u = sum(u)
        v = sum(v)
    else:
        u = -strength / (2 * math.pi) * ((X - xd) ** 2 - (Y - yd) ** 2) / (
            (X - xd) ** 2 + (Y - yd) ** 2) ** 2
        v = -strength / (2 * math.pi) * 2 * (X - xd) * (Y - yd) / ((X - xd) **
            2 + (Y - yd) ** 2) ** 2
    return u, v


def get_stream_function_doublet(strength, xd, yd, X, Y):
    """
    Returns the stream-function generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    psi = -strength / (2 * math.pi) * (Y - yd) / ((X - xd) ** 2 + (Y - yd) ** 2
        )
    return psi


<mask token>
pyplot.figure(figsize=(size, (y_end - y_start) / (x_end - x_start) * size))
pyplot.xlabel('x', fontsize=16)
pyplot.ylabel('y', fontsize=16)
pyplot.xlim(x_start, x_end)
pyplot.ylim(y_start, y_end)
pyplot.streamplot(X, Y, u, v, density=2, linewidth=1, arrowsize=2,
    arrowstyle='->')
pyplot.scatter(x_source, y_source, s=10, color='#CD2305', marker='o',
    linewidth=0)
pyplot.contour(X, Y, psi, levels=[1, 3], colors='#CD2305', linewidths=2,
    linestyles='solid')
pyplot.show()

<mask token>


def velocity_grid(strength, x_source, y_source, X, Y):
    """
    Parameters
    ----------
    strength: list[floats] or float
        Strength of the sources/sinks
    x_source: list[floats] or float
        x-coordinates of the sources and sinks.
    y_source: list[floats] or float
        y-coordinates of the sources and sinks.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.

    Returns
    -------
    u_source 2D Numpy array of floats
        x-velocity field
    v_source 2D Numpy array of floats
        y-velocity field

    Note
    ----
    sources are superimposed into one velocity field
    """
    if isinstance(x_source, list):
        u_source = [0.0] * len(x_source)
        v_source = [0.0] * len(y_source)
        for i in range(0, len(x_source)):
            u_source[i] = strength[i] / (2 * math.pi) * (X - x_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
            v_source[i] = strength[i] / (2 * math.pi) * (Y - y_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
        u_source = sum(u_source)
        v_source = sum(v_source)
    else:
        u_source = strength / (2 * math.pi) * (X - x_source) / ((X -
            x_source) ** 2 + (Y - y_source) ** 2)
        v_source = strength / (2 * math.pi) * (Y - y_source) / ((X -
            x_source) ** 2 + (Y - y_source) ** 2)
    return u_source, v_source


def get_stream_function(strength, x_source, y_source, X, Y):
    """
    Returns the stream-function generated by a source/sink.
    
    Parameters
    ----------
    strength: float
        Strength of the source/sink.
    xs: float
        x-coordinate of the source (or sink).
    ys: float
        y-coordinate of the source (or sink).
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    if isinstance(x_source, list):
        psi = [0.0] * len(x_source)
        for i in range(0, len(x_source)):
            psi[i] = strength[i] / (2 * numpy.pi) * numpy.arctan2(Y -
                y_source[i], X - x_source[i])
        psi = sum(psi)
    else:
        psi = strength / (2 * numpy.pi) * numpy.arctan2(Y - y_source, X -
            x_source)
    return psi


def get_velocity_doublet(strength, xd, yd, X, Y):
    """
    Returns the velocity field generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    u: 2D Numpy array of floats
        x-component of the velocity vector field.
    v: 2D Numpy array of floats
        y-component of the velocity vector field.
    """
    if isinstance(xd, list):
        u = [0.0] * len(xd)
        v = [0.0] * len(yd)
        for i in range(0, len(xd)):
            u[i] = -strength[i] / (2 * math.pi) * ((X - xd[i]) ** 2 - (Y -
                yd[i]) ** 2) / ((X - xd[i]) ** 2 + (Y - yd[i]) ** 2) ** 2
            v[i] = -strength[i] / (2 * math.pi) * 2 * (X - xd[i]) * (Y - yd[i]
                ) / ((X - xd[i]) ** 2 + (Y - yd[i]) ** 2) ** 2
        u = sum(u)
        v = sum(v)
    else:
        u = -strength / (2 * math.pi) * ((X - xd) ** 2 - (Y - yd) ** 2) / (
            (X - xd) ** 2 + (Y - yd) ** 2) ** 2
        v = -strength / (2 * math.pi) * 2 * (X - xd) * (Y - yd) / ((X - xd) **
            2 + (Y - yd) ** 2) ** 2
    return u, v


def get_stream_function_doublet(strength, xd, yd, X, Y):
    """
    Returns the stream-function generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    psi = -strength / (2 * math.pi) * (Y - yd) / ((X - xd) ** 2 + (Y - yd) ** 2
        )
    return psi


N = 50
x_start, x_end = 0.0, 3.0
y_start, y_end = 0.0, 2.0
x = numpy.linspace(x_start, x_end, N)
y = numpy.linspace(y_start, y_end, N)
X, Y = numpy.meshgrid(x, y)
<mask token>
strength = [5.0, -1.0]
x_source = [1.0, 1.5]
y_source = [1.0, 1.0]
u_inf = 2.0
u_freestream = u_inf * numpy.ones((N, N), dtype=float)
v_freestream = numpy.zeros((N, N), dtype=float)
psi_freestream = u_inf * Y
u_source, v_source = velocity_grid(strength, x_source, y_source, X, Y)
psi_source = get_stream_function(strength, x_source, y_source, X, Y)
u = u_freestream + u_source
v = v_freestream + v_source
psi = psi_freestream + psi_source
size = 10
pyplot.figure(figsize=(size, (y_end - y_start) / (x_end - x_start) * size))
pyplot.xlabel('x', fontsize=16)
pyplot.ylabel('y', fontsize=16)
pyplot.xlim(x_start, x_end)
pyplot.ylim(y_start, y_end)
pyplot.streamplot(X, Y, u, v, density=2, linewidth=1, arrowsize=2,
    arrowstyle='->')
pyplot.scatter(x_source, y_source, s=10, color='#CD2305', marker='o',
    linewidth=0)
pyplot.contour(X, Y, psi, levels=[1, 3], colors='#CD2305', linewidths=2,
    linestyles='solid')
pyplot.show()

import math
import numpy
from matplotlib import pyplot


def velocity_grid(strength, x_source, y_source, X, Y):
    """
    Parameters
    ----------
    strength: list[floats] or float
        Strength of the sources/sinks
    x_source: list[floats] or float
        x-coordinates of the sources and sinks.
    y_source: list[floats] or float
        y-coordinates of the sources and sinks.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.

    Returns
    -------
    u_source 2D Numpy array of floats
        x-velocity field
    v_source 2D Numpy array of floats
        y-velocity field

    Note
    ----
    sources are superimposed into one velocity field
    """
    if isinstance(x_source, list):
        u_source = [0.0] * len(x_source)
        v_source = [0.0] * len(y_source)
        for i in range(0, len(x_source)):
            u_source[i] = strength[i] / (2 * math.pi) * (X - x_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
            v_source[i] = strength[i] / (2 * math.pi) * (Y - y_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
        u_source = sum(u_source)
        v_source = sum(v_source)
    else:
        u_source = strength / (2 * math.pi) * (X - x_source) / ((X -
            x_source) ** 2 + (Y - y_source) ** 2)
        v_source = strength / (2 * math.pi) * (Y - y_source) / ((X -
            x_source) ** 2 + (Y - y_source) ** 2)
    return u_source, v_source


def get_stream_function(strength, x_source, y_source, X, Y):
    """
    Returns the stream-function generated by a source/sink.
    
    Parameters
    ----------
    strength: float
        Strength of the source/sink.
    xs: float
        x-coordinate of the source (or sink).
    ys: float
        y-coordinate of the source (or sink).
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    if isinstance(x_source, list):
        psi = [0.0] * len(x_source)
        for i in range(0, len(x_source)):
            psi[i] = strength[i] / (2 * numpy.pi) * numpy.arctan2(Y -
                y_source[i], X - x_source[i])
        psi = sum(psi)
    else:
        psi = strength / (2 * numpy.pi) * numpy.arctan2(Y - y_source, X -
            x_source)
    return psi


def get_velocity_doublet(strength, xd, yd, X, Y):
    """
    Returns the velocity field generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    u: 2D Numpy array of floats
        x-component of the velocity vector field.
    v: 2D Numpy array of floats
        y-component of the velocity vector field.
    """
    if isinstance(xd, list):
        u = [0.0] * len(xd)
        v = [0.0] * len(yd)
        for i in range(0, len(xd)):
            u[i] = -strength[i] / (2 * math.pi) * ((X - xd[i]) ** 2 - (Y -
                yd[i]) ** 2) / ((X - xd[i]) ** 2 + (Y - yd[i]) ** 2) ** 2
            v[i] = -strength[i] / (2 * math.pi) * 2 * (X - xd[i]) * (Y - yd[i]
                ) / ((X - xd[i]) ** 2 + (Y - yd[i]) ** 2) ** 2
        u = sum(u)
        v = sum(v)
    else:
        u = -strength / (2 * math.pi) * ((X - xd) ** 2 - (Y - yd) ** 2) / (
            (X - xd) ** 2 + (Y - yd) ** 2) ** 2
        v = -strength / (2 * math.pi) * 2 * (X - xd) * (Y - yd) / ((X - xd) **
            2 + (Y - yd) ** 2) ** 2
    return u, v


def get_stream_function_doublet(strength, xd, yd, X, Y):
    """
    Returns the stream-function generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    psi = -strength / (2 * math.pi) * (Y - yd) / ((X - xd) ** 2 + (Y - yd) ** 2
        )
    return psi


N = 50
x_start, x_end = 0.0, 3.0
y_start, y_end = 0.0, 2.0
x = numpy.linspace(x_start, x_end, N)
y = numpy.linspace(y_start, y_end, N)
X, Y = numpy.meshgrid(x, y)
<mask token>
strength = [5.0, -1.0]
x_source = [1.0, 1.5]
y_source = [1.0, 1.0]
u_inf = 2.0
u_freestream = u_inf * numpy.ones((N, N), dtype=float)
v_freestream = numpy.zeros((N, N), dtype=float)
psi_freestream = u_inf * Y
u_source, v_source = velocity_grid(strength, x_source, y_source, X, Y)
psi_source = get_stream_function(strength, x_source, y_source, X, Y)
u = u_freestream + u_source
v = v_freestream + v_source
psi = psi_freestream + psi_source
size = 10
pyplot.figure(figsize=(size, (y_end - y_start) / (x_end - x_start) * size))
pyplot.xlabel('x', fontsize=16)
pyplot.ylabel('y', fontsize=16)
pyplot.xlim(x_start, x_end)
pyplot.ylim(y_start, y_end)
pyplot.streamplot(X, Y, u, v, density=2, linewidth=1, arrowsize=2,
    arrowstyle='->')
pyplot.scatter(x_source, y_source, s=10, color='#CD2305', marker='o',
    linewidth=0)
pyplot.contour(X, Y, psi, levels=[1, 3], colors='#CD2305', linewidths=2,
    linestyles='solid')
pyplot.show()

import math
import numpy
from matplotlib import pyplot


def velocity_grid(strength, x_source, y_source, X, Y):
    """
    Parameters
    ----------
    strength: list[floats] or float
        Strength of the sources/sinks
    x_source: list[floats] or float
        x-coordinates of the sources and sinks.
    y_source: list[floats] or float
        y-coordinates of the sources and sinks.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.

    Returns
    -------
    u_source 2D Numpy array of floats
        x-velocity field
    v_source 2D Numpy array of floats
        y-velocity field

    Note
    ----
    sources are superimposed into one velocity field
    """
    if isinstance(x_source, list):
        u_source = [0.0] * len(x_source)
        v_source = [0.0] * len(y_source)
        for i in range(0, len(x_source)):
            u_source[i] = strength[i] / (2 * math.pi) * (X - x_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
            v_source[i] = strength[i] / (2 * math.pi) * (Y - y_source[i]) / (
                (X - x_source[i]) ** 2 + (Y - y_source[i]) ** 2)
        u_source = sum(u_source)
        v_source = sum(v_source)
    else:
        u_source = strength / (2 * math.pi) * (X - x_source) / ((X - x_source) ** 2 + (Y - y_source) ** 2)
        v_source = strength / (2 * math.pi) * (Y - y_source) / ((X - x_source) ** 2 + (Y - y_source) ** 2)
    return u_source, v_source


def get_stream_function(strength, x_source, y_source, X, Y):
    """
    Returns the stream-function generated by a source/sink.
    
    Parameters
    ----------
    strength: float
        Strength of the source/sink.
    xs: float
        x-coordinate of the source (or sink).
    ys: float
        y-coordinate of the source (or sink).
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    if isinstance(x_source, list):
        psi = [0.0] * len(x_source)
        for i in range(0, len(x_source)):
            psi[i] = strength[i] / (2 * numpy.pi) * numpy.arctan2((Y - y_source[i]), (X - x_source[i]))
        psi = sum(psi)
    else:
        psi = strength / (2 * numpy.pi) * numpy.arctan2((Y - y_source), (X - x_source))

    return psi

def get_velocity_doublet(strength, xd, yd, X, Y):
    """
    Returns the velocity field generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    u: 2D Numpy array of floats
        x-component of the velocity vector field.
    v: 2D Numpy array of floats
        y-component of the velocity vector field.
    """
    if isinstance(xd, list):
        u = [0.0] * len(xd)
        v = [0.0] * len(yd)
        for i in range(0, len(xd)):
            u[i] = - strength[i] / (2 * math.pi) * ((X - xd[i]) ** 2 - (Y - yd[i]) ** 2) / ((X - xd[i]) ** 2 + (
                Y - yd[i]) ** 2) ** 2
            v[i] = - strength[i] / (2 * math.pi) * 2 * (X - xd[i]) * (Y - yd[i]) / ((X - xd[i]) ** 2 + (
                Y - yd[i]) ** 2) ** 2
        u = sum(u)
        v = sum(v)
    else:
        u = - strength / (2 * math.pi) * ((X - xd) ** 2 - (Y - yd) ** 2) / ((X - xd) ** 2 + (Y - yd) ** 2) ** 2
        v = - strength / (2 * math.pi) * 2 * (X - xd) * (Y - yd) / ((X - xd) ** 2 + (Y - yd) ** 2) ** 2

    return u, v

def get_stream_function_doublet(strength, xd, yd, X, Y):
    """
    Returns the stream-function generated by a doublet.
    
    Parameters
    ----------
    strength: float
        Strength of the doublet.
    xd: float
        x-coordinate of the doublet.
    yd: float
        y-coordinate of the doublet.
    X: 2D Numpy array of floats
        x-coordinate of the mesh points.
    Y: 2D Numpy array of floats
        y-coordinate of the mesh points.
    
    Returns
    -------
    psi: 2D Numpy array of floats
        The stream-function.
    """
    psi = - strength / (2 * math.pi) * (Y - yd) / ((X - xd) ** 2 + (Y - yd) ** 2)

    return psi

# establish grid
N = 50
x_start, x_end = 0.0, 3.0
y_start, y_end = 0.0, 2.0
x = numpy.linspace(x_start, x_end, N)
y = numpy.linspace(y_start, y_end, N)
# mesh creation
X, Y = numpy.meshgrid(x, y)

# source + sink input information
"""
N_source = 100
strength = [0.1] * N_source # source strength
x_source = list(numpy.linspace(-0.5, 1, N_source)) #x_coordinates
y_source = list(numpy.linspace(0.5, 0.0, N_source)) 
#y_source = [0.0] * N_source # y_coordinates
"""
strength = [5.0, -1.0]
x_source = [1.0, 1.5]
y_source = [1.0, 1.0]

# adding in free stream velocity field
u_inf = 2.0

u_freestream = u_inf * numpy.ones((N, N), dtype=float)
v_freestream = numpy.zeros((N, N), dtype=float)

# compute the stream function
psi_freestream = u_inf * Y

# compute velocity grid
u_source, v_source = velocity_grid(strength, x_source, y_source, X, Y)
psi_source = get_stream_function(strength, x_source, y_source, X, Y)

# compute final velocity grid
u = u_freestream + u_source
v = v_freestream + v_source
psi = psi_freestream + psi_source


# plot the gird of points
size = 10
pyplot.figure(figsize=(size, (y_end-y_start)/(x_end-x_start)*size))
pyplot.xlabel('x', fontsize = 16)
pyplot.ylabel('y', fontsize = 16)
pyplot.xlim(x_start, x_end)
pyplot.ylim(y_start, y_end)
pyplot.streamplot(X, Y, u, v, density=2, linewidth=1, arrowsize=2, arrowstyle='->')
pyplot.scatter(x_source, y_source, s=10, color = '#CD2305', marker = 'o', linewidth = 0)
pyplot.contour(X,Y,psi,
               #levels=[-strength[0]/2.,0, +strength[0]/2.],
               levels=[1,3],
               colors='#CD2305',linewidths=2, linestyles='solid')
pyplot.show()
