from IPython.display import display, Image
from IPython.display import display_pretty, display_html, display_jpeg, display_png, display_json, display_latex, display_svg
from IPython.display import Latex
from IPython.display import HTML
Image(filename='logo3.jpg', width=100)

Image(filename='spacing.jpg', width=400)

# U = internal structure + Ubc = boudary conditions
rhs = np.ravel((U+Ubc))

from scipy import sparse
import numpy as np

def Lap(size, delta, boudary_cond='Dirichlet'):

    """Function returning Matrix representing Laplace operator"""
    """arguments:"""
    """size - size of matrix"""
    """delta - discretisation step"""
    """boundary_cond - type of boundary_condition """
    
    
    Lap_row = np.concatenate((np.array([-1, 2, -1]), np.zeros(size-2)))
    Lap_op = Lap_row[:]

    for i in range(size-1):
        Lap_op = np.concatenate((Lap_op, Lap_row))

    deletion = [0] + list(range(-size+1, 0))

    for i in deletion:
        Lap_op = np.delete(Lap_op, i)

    look_up_tab = {'Neumann': 1, 'Dirichlet': 2, 'Dirichlet_between': 3}
    Lap_op[0] = Lap_op[-1] = look_up_tab[boudary_cond]

    return np.reshape(Lap_op, (size, size)) / delta ** 2


nx = 3
ny = 3
Re = 1
dt = 1
hx = 1
hy = 1

Lu = sparse.eye(nx*(ny-1))+dt / Re*(sparse.kron(sparse.eye(ny), Lap(nx-1, hx)) + 
                                    sparse.kron(Lap(ny, hy), sparse.eye(nx-1)))
Lv = sparse.eye(nx*(ny-1))+dt / Re*(sparse.kron(sparse.eye(ny-1), Lap(nx, hx)) +
                                    sparse.kron(Lap(ny-1, hy), sparse.eye(nx)))
print(Lu)

import numpy as np
import math as mt

def avg(A, k=1):
"""
    Returns horizontally averaged matrix
"""
    if A.shape[0] == 1:
        A = np.transpose(A)
    if k < 2:
        B = (A[1::][:] + A[0:-1][:])/2
    else:
        B = avg(A, k-1)
    if len(A.shape) < 2:
        B = np.reshape(B, (1, len(B)))
    return B

#Copies of the field:
Ue = U
Ve = V

#Treating U^2, V^2 terms
Ua = avg(Ue[:, 1:-1])   #We are leaving boundary conditions untouched thus [1:-1]
Ud = np.diff(Ue[:, 1:-1], axis=0)/2.0
Va = avg(Ve[1:-1][:].T).T
Vd = np.diff(Ve[1:-1][:])/2.0

U2x = np.diff(Ua**2, axis=0) / hx
V2y = np.diff(Va**2) / hy

#UV Terms
Ua = avg(Ue.T)   #vertically averaged
Ud = np.diff(Ue)/2.0
Va = avg(Ve)    #horrizontally averaged
Vd = np.diff(Ve, axis=0)/2.0

UVx = np.diff(Ua * Va, axis=0) / hx
UVy = np.diff(Ua * Va) / hy

#Update equation:
U += - dt * (UVy[1:-1, :] + U2x)
V += - dt * (UVx[:, 1:-1] + V2y)

class Operators:
    """
        Class responsible for shorthand notation
        Functions:
        cncat - concatenation ->
            matrix_tuple -> tuple of matrixes to concatenate
            dim -> vertical/horizontal direction, vertical by default
        vec2col -> transforming np array generated as (x,) to column vector (x,1)
        vec2row -> transforming np array generated as (x,) to row vector (1,x)
    """
    def cncat(self, matrix_tuple, dim='v'):
        """Concatenation"""
        if dim == 'v':
            return np.concatenate(matrix_tuple, axis=0)
        if dim == 'h':
            return np.concatenate(matrix_tuple, axis=1)
    def vec2col(self, vec):
        return np.reshape(vec, (max(vec.shape), 1))

    def vec2row(self, vec):
        return np.reshape(vec, (1, max(vec.shape)))


from IPython.display import HTML
HTML('<iframe src=http://by2pie.wordpress.com width=900 height=500></iframe>')