# solves dirichlet problem with -u''(x)=f(x), 0<x<1, u(0)=a1, u(1)=a2
# user specified functions u(x) and grid x as input to code

import numpy as np
import matplotlib.pyplot as plt

# 
def f(x):
    return (np.pi**2)*np.cos(np.pi*x)

def U_inhom_dir(x,a1,a2):
    n = x.shape[0]
    h = np.diff(x)
    A = np.zeros((n,n))
    b = np.zeros((n,1))
    # loop over elements
    for i in range(n-1):
        A[i:i+2][:,i:i+2] = A[i:i+2][:,i:i+2] + np.array([[1,-1],[-1,1]])/h[i]
        # use midpoint ('barycentric') one node quadrature
        b[i:i+2] = b[i:i+2] + f(np.mean(x[i:i+2]))*h[i]/2
    # construct lifting
    R = np.vstack([a1,np.zeros((n-2,1)),a2])
    b = b - A.dot(R)
    # implement boundary conditions
    A[0] = np.zeros((1,n))
    A[-1] = np.zeros((1,n))
    A[0][0] = 1
    A[-1][-1] = 1
    b[0] = 0
    b[-1] = 0
    # solve homogeneous problem
    u = np.linalg.solve(A,b)
    # add lifting back
    u[0] = a1
    u[-1] = a2
    return u

a1=1;
a2=-1;
x = np.linspace(0,1,100)
u=U_inhom_dir(x,a1,a2)   
plt.plot(x,u)