dx = 0.3
x = np.arange(0, 10, dx) # returns [0, dx, 2dx, 3dx, 4dx, 5dx, ...]
print(x)
f1 = np.sin(x)
f2 = x**2/100
f3 = np.log(1+x)-1
fs = [f1, f2, f3]

for i in range(3): plt.plot(x, fs[i])

df1 = np.cos(x)
df2 = x/50
df3 = 1/(1+x)
dfs = [df1, df2, df3]

"""
Explanation: Numerical differential equations
In the simplest methods for solving differential equations numerically, one goes back to the definition of the differential operator: lim dx-&gt;0 (f(t + dx) - f(t))/dx and simply take dx small, but finite, which will approximate the true derivative.
We start by understanding this.
Numerical differentiation
We start by defining a few functions for which we know their analytical derivative.
End of explanation
"""


def derivative(f, dx):
    return (f[1:] - f[:-1])/dx # returned function is one value shorter than input f
ndfs = [derivative(f, dx) for f in fs]

for i in range(3): plt.plot(x[:-1], ndfs[i], lw=3, alpha=0.5)
for i in range(3): plt.plot(x, dfs[i], '--k')
plt.show()

"""
Explanation: Now let us see if we can calculate these derivatives numerically.
End of explanation
"""


def central_derivative(f, dx):
    return (f[2:] - f[:-2])/(2*dx) # returned function is two values shorter than input f
ndfs = [central_derivative(f, dx) for f in fs]
for i in range(3): plt.plot(x[1:-1], ndfs[i], lw=3, alpha=0.5)
for i in range(3): plt.plot(x, dfs[i], '--k')
plt.show()

"""
Explanation: That worked pretty well, but we can do even better by using central differences for estimating the deriavtive.
End of explanation
"""


def analytical_solution(t, k, d, x0, v0):
    d += 0j # Exploiting complex numbers for one general solution
    return np.real((np.exp(-((d*t)/2))*(np.sqrt(d**2-4*k)*x0*np.cosh(1/2*np.sqrt(d**2-4*k)*t)+(2*v0+d* \
           x0)*np.sinh(1/2*np.sqrt(d**2-4*k)*t)))/np.sqrt(d**2-4*k))

dt = 0.02
t = np.arange(0,25,dt)
k = 5; d = 0.3; x0 = 0; v0 = 1

x = np.zeros_like(t)
v = np.zeros_like(t)
x[0] = x0; v[0] = v0

for i in range(len(t)-1): # Step through time with step size dt
    v[i + 1] = v[i] + (-k * x[i] - d * v[i]) * dt
    x[i + 1] = x[i] + v[i+1] * dt # v[i+1] makes a big stability difference!

plt.plot(t, x, lw=3, alpha=0.5)
plt.plot(t, analytical_solution(t, k, d, x0, v0), '--k')
plt.show()

"""
Explanation: The error is order dx for left/right derivatives, but only dx^2 for central derivatives.
Initial Value Problems
We now turn to simple ODE's. Consider the damped harmonic oscillator equation
x''(t) = - k x(t) - d x'(t)
which we can write as
x'(t) = v(t)
v'(t) = -k x(t) - d v(t).
We want to solve this with initial conditions x(0) = x0, v(0) = v0.
End of explanation
"""


dx = 0.01
x = np.arange(0, 1, dx)
n = len(x)
y0 = 1
y1 = 10

"""
Explanation: Exercise
(Skip this exercise if you feel confident it's easy)
Change the above code to solve 
x''(t) = - k sin(x(t))
If time permits: lookup documentation for scipy.integrate.odeint and use that to solve the above equation. This will be much more efficient. But do other exercises first.
Boundary Value Problems
Now we want to solve the equation
y''(x) = y(x)
y(0) = y0
y(1) = y1
We can no longer iterate forwards from a given start point as we could before, where both x(0) and x'(0) were known.
There are many methods to solve systems as the above (shootings methods, relaxation methods, etc), but we will use the simplest direct method here. Again we base it on finite differences.
End of explanation
"""


y = x**4 # not correct

"""
Explanation: Imagine that we had calculated y. It would be some array:
End of explanation
"""


L = np.diag(np.ones(n - 1), 1) + np.diag(np.ones(n - 1), -1) - 2 * np.diag(np.ones(n))
print(L)
L = L/dx**2

ypp = np.dot(L, y) # matrix product L y
plt.plot(x, ypp, lw=3, alpha=0.5)
plt.plot(x, 4*3*x**2, '--k')
plt.axis([0, 1, -3, 15])
plt.show()

"""
Explanation: How would we calculate the vector y''? In particular, can we find a matrix L such that y'' = L y?
Double derivatives can be estimated by finite differences as (see https://en.wikipedia.org/wiki/Finite_difference_coefficient)
y''(x) = (y(x + dx) + y(x - dx) - 2 y(x))/dx^2
Let's build a matrix based on that:
End of explanation
"""


I = np.eye(n)
A = L - I

"""
Explanation: All looks good, expect the endpoints. This was expected. This is why we have boundary conditions.
We can write our equation y''(x) = y(x) as y''(x) - y(x) = 0 and by introducing the double derivative operator L and the identity operator I, we could also write it as
(L - I) y = 0
The idenitity operator for vectors is of course just the idenitity matrix
End of explanation
"""


A[0, :] = 0
A[0, 0] = 1
A[-1, :] = 0
A[-1, -1] = 1
b = np.zeros(n)
b[0] = y0
b[-1] = y1

np.set_printoptions(suppress=True)
print(np.round(A))
print('* y == ')
print(b)

"""
Explanation: Where A is defining the full problem.
The problem is still not well-defined as we need to somehow implement our boundary conditions.
Let's do it, and explain afterwards:
End of explanation
"""


from numpy.linalg import solve
y = solve(A, b)

analytical = (np.exp(-x)*(np.exp(1)*(np.exp(1)*y0-y1)+np.exp(2*x)*(-y0+np.exp(1)*y1)))/(-1+np.exp(2))

plt.plot(x, y, lw=3, alpha=0.5)
plt.plot(x, analytical, '--k')
plt.show()

"""
Explanation: The first row extracts y[0] and the equation then sets it equal to y0(=1) and the last row extracts y[-1] and sets it equal to y1(=10).
All other lines sets double derivatives equal to zero.
Solving the above linear system of equations is easy.
Numpy linear algebra module can do it for us:
End of explanation
"""


dx = 0.01
x = np.arange(-np.pi, np.pi, dx)
y = 1/4*np.exp(-np.pi)/np.sinh(np.pi)*(-2+np.exp(x)*(1+2*np.exp(np.pi)+np.cos(x)+ \
        np.sin(x)-np.exp(2*(np.pi))*(1+np.cos(x)+np.sin(x))))
plt.plot(x, y, '--k')
plt.axis([-np.pi, np.pi, -6, 1.5])
plt.show()

"""
Explanation: It is great for testing one's code that we can compare to analytical solutions, but of course the strength of numerical methods is that we can solve equations that we do not have analytical solutions for.
Exercise
Solve
y''(x) = y'(x) + sin(x)*exp(x)
y(-pi) = 0
y(pi) = 1
The analytical solution is given below.
End of explanation
"""


dx = 0.01
x = np.arange(0, 7, dx)
y = np.arange(-3, 3, dx)
X, Y = np.meshgrid(x, y)
U = np.exp(X) * np.sin(Y)

plt.imshow(U, extent=(min(x), max(x), max(y), min(y)))
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar()
plt.show()

"""
Explanation: Extra exercise
What about Neumann boundary conditions? Try writing a solver for
y''(x) = y(x) + sin(x)*exp(x)
y(-pi) = 1
y'(pi) = 0
As a check, the correct solution has y(0) = -0.956331.
Partial differential equations
Just as PDE's are harder to solve than ODE's, so are PDE's much harder to implement numerically.
Most courses tend to avoid these and say they "generalise from ODE's", which I would argue is far from true.
This being a maths department and PDE solving might be relevant, so we will cover this area. But focus on simple methods.
3D functions
We start by showing how to plot
u(x,y) = exp(x) sin(y)
End of explanation
"""


print(X.shape, Y.shape, U.shape)
print(X)

"""
Explanation: In the above np.meshgrid has made the 2D arrays X,Y:
End of explanation
"""


dx = 0.02
x = np.arange(0, 1 + dx, dx)
m = len(x)
print(x[0], x[-1])
X, Y = np.meshgrid(x, x)
shape = X.shape

"""
Explanation: Example problem
Let us solve the steady state heat (Laplace) equation with known temperature boundary conditions
Laplacian u(x, y) = 0
u(x, 0) = 1
u(0, y) = 1 - y
u(x, 1) = 0
y(1, y) = 0
We start by defining our domain:
End of explanation
"""


def to_vector(mat):
    return np.ravel(mat)
def to_matrix(vec):
    return np.reshape(vec, shape)

print(X.shape, '=>', to_vector(X).shape)
print((X == to_matrix(to_vector(X))).all())
print((Y == to_matrix(to_vector(Y))).all())

"""
Explanation: We wish to solve for U, but currently U, like X and Y, would be a matrix. Linear systems tend to work on vectors, so we should formulate our system such that we can solve for a vector.
For this purpose we define functions that take us from vector to matrix and vice versa.
End of explanation
"""


x = to_vector(X)
y = to_vector(Y)
n = len(x)

"""
Explanation: So we define 1D arrays of our coordinates:
End of explanation
"""


L = np.zeros((n, n))
for i in range(n):
    L[i,i] = -4
    
    j = np.argmin( (x[i] + dx - x)**2 + (y[i] - y)**2 ) # Find index j in vectors for point (x[i]+dx, y[i])
    if i!=j: L[i,j] = 1 # If i==j, we are at the boundary of the domain
    
    j = np.argmin( (x[i] - dx - x)**2 + (y[i] - y)**2 ) # Find index j in vectors for point (x[i]-dx, y[i])
    if i!=j: L[i,j] = 1
    
    j = np.argmin( (x[i] - x)**2 + (y[i] + dx - y)**2 ) # Find index j in vectors for point (x[i], y[i]+dx)
    if i!=j: L[i,j] = 1
    
    j = np.argmin( (x[i] - x)**2 + (y[i] - dx - y)**2 ) # Find index j in vectors for point (x[i], y[i]-dx)
    if i!=j: L[i,j] = 1
print(L)
L = L/dx**2

"""
Explanation: The discrete Laplacian looks like this:
0   1   0
1  -4   1
0   1   0
divided by dx^2
There are smart ways to make this, but it depends on how you form the vectors.
We should not dwelve too long at the details of how we've done it here (ravel, reshape), but instead we will build the Laplacian in a silly way using a loop. This way will work no matter how scrambled the coordinates are in the vector.
End of explanation
"""


L_quick = -4 * np.eye(n) + np.diag(np.ones(n-m), m) + np.diag(np.ones(n-m), -m)
a = np.ones(n-1); a[(m-1)::m] = 0
L_quick += np.diag(a,1) + np.diag(a,-1)
L_quick = L_quick/dx**2

print( (L == L_quick).all() )

"""
Explanation: As mentioned, this can be done much more efficiently (and without using loops).
But this requires that one understands how the reorganisation to vectors work.
Anyway, just to show how one could do it:
End of explanation
"""


b = np.zeros(n)
for i in range(n):
    if (x[i]==0 or x[i]==1 or y[i]==0 or y[i]==1): # For any boundary point
        L[i, :] = 0
        L[i, i] = 1
    
    # BC points that are not equal to zero:
    if x[i] == 0:
        b[i] = 1 - y[i]
    elif y[i] == 0:
        b[i] = 1

"""
Explanation: Now we implement the boundary conditions:
End of explanation
"""


from scipy.linalg import solve
u = solve(L, b)

U = to_matrix(u)
plt.imshow(U, extent=(min(x), max(x), max(y), min(y)))
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar()
plt.title('Temperature distriubtion of plate')
plt.show()

"""
Explanation: And lastly, we solve:
End of explanation
"""


