Til emil

import numpy as np
import matplotlib.pyplot as plt
try:
    from pprint import pprint
except ModuleNotFoundError:
    pass

def cholesky(A):
    '''
    A er en positiv definitt matrise som returnerer en nedretriangulær matrise L slik at A=L L^T
    '''
    assert A.shape[0] == A.shape[1], "The inupt matrix is not square"
    n = A.shape[0]
    L = np.zeros_like(A)

    for i in range(n):
        for j in range(i):
            L[i, j] = (A[i, j] - np.sum(L[i,:j]*L[j,:j]) ) / L[j, j]
        L[i, i] = np.sqrt(A[i, i] - np.sum(L[i,:i]**2))

    assert np.allclose(L, np.linalg.cholesky(A)), f"Cholesky failed: L:\n{L}\n" + \
                                                   f"np:\n{np.linalg.cholesky(A)}\nL-np:\n{L-np.linalg.cholesky(A)}"

    # print((L@L.T - A))
    # return np.linalg.cholesky(A)
    return L

# ----------------------------------------------------------------
#
# ----------------------------------------------------------------

def substit(L,b):
    '''For en nedretriangulær matrise L (nxn) og en vektor b (nx1) finn først c (nx1) slik at Lc=b
    og deretter x (nx1) slik at Lx=c'''
    assert (b.size == L.shape[0] == L.shape[1]) if b.ndim == 1 else True
    n = b.shape[0]

    c = np.empty_like(b)

    for i in range(n):
        c[i] = ( b[i] - np.sum(L[i, :i] * c[:i])) / L[i, i]

    assert np.allclose(L @ c, b), f"Forwards substitution failed:\n{L @ c - b}"

    x = np.empty_like(b)
    for i in range(n-1, -1, -1):
        x[i] = (c[i] - np.sum(L.T[i,i+1:]*x[i+1:])) / L[i, i]

    assert np.allclose(L.T @ x, c), f"Backwards substitution failed:\n{L.T @ x - c}"

    return x

import scipy.linalg as sclin
def catch_substit(L, b, skip=False):
    if not skip:
        try:
            substit(L, b)
            return x
        except AssertionError as error:
            #print(error)
            print("AssertionError in substit. Solving with scipy.linalg:")

    c = sclin.solve_triangular(L, b)
    x = sclin.solve_triangular(L.T, c)
    return x

# ----------------------------------------------------------------
# Test once:
# ----------------------------------------------------------------

# Verifikasjon oppgave a og b
A = np.array([[1,2,3],[2,5,4],[3,4,14]])
b = np.array([-2,-8,3])

# Kall dine funksjoner cholesky og substit her
L = cholesky(A)
x = substit(L, b)

# Skriv ut L og x, (fjern # eller modifiser flg 2 linjer)
print('L=\n',L)
print('x=\n',x)




# ----------------------------------------------------------------
# Define helper functions (given in task)
# ----------------------------------------------------------------

def genM(n,Delta_t):
    Delta_x = 1./(n+1)
    r=Delta_t/Delta_x**2
    ee=np.ones((n,))
    B=(1+4*r)*np.diag(ee)-r*np.diag(ee[1:],-1)-r*np.diag(ee[1:],1)
    In=np.diag(ee)
    Fn=np.diag(ee[1:],1)
    Gn=np.diag(ee[1:],-1)
    M=np.kron(In,B)-r*np.kron(Fn,In)-r*np.kron(Gn,In)
    return M


def genU0(n):
    p=n**2/4
    Z0=np.zeros((n,n))
    Z0[n-1,0]=p
    Z0[0,n-1]=p
    Z0[0,0] = p
    Z0[n-1,n-1]=p
    U0 = np.reshape(Z0,(n**2,1))
    return U0

def surfplot(U):
    from mpl_toolkits.mplot3d import Axes3D
    from matplotlib import cm
    n2=U.shape[0]
    n=int(np.sqrt(n2))
    if n2 != n**2:
        print('Antall elementer i U må være et kvadrattall\n')
        return
    Delta_x=1/(n+1)
    xx=np.linspace(Delta_x,1-Delta_x,n)
    yy=np.linspace(Delta_x,1-Delta_x,n)
    X,Y = np.meshgrid(xx,yy)
    Z = np.reshape(U,(n,n))
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    surf = ax.plot_surface(X, Y, Z, cmap=cm.plasma,
                       linewidth=1, antialiased=False)
    plt.show()
    return surf



# ----------------------------------------------------------------
#
# ----------------------------------------------------------------

# Her skriver du kode for å løse oppgave c
#
#  1. Hent U0 og beregn U1, U2, U3 ved bruk av dine cholesky- og substitfunksjoner
#  2. Plott U0 og U3 med surfplot-funksjonen

n, dt = 30, 0.01

M = genM(n, dt)
U_0 = genU0(n)
U = [None for n in range(4)]
U[0] = U_0
surfplot(U[0])


L = cholesky(M)

for i in range(len(U)-1):
    U[i+1] = catch_substit(L, U[i], skip=True)


# print(U_0)
# print(U_1)
surfplot(U[3])

# For å svare på Kontrollspørsmål 2, aktiver følgende
Z=np.reshape(U[3],(n,n))
print(Z[14,15])


# ----------------------------------------------------------------
#
# ----------------------------------------------------------------



# ----------------------------------------------------------------
#
# ----------------------------------------------------------------