#!/usr/bin/python3
# Cooling of a homogeneous ball

import numpy as np
import matplotlib.pyplot as plt
import scipy.linalg as la

D = 4.E-6    # diffusion coefficient
c = 20.      # cooling constant
R = 5.E-2    # ball radius
Ti = 373.    # temperature at t=0
Tw = 290.    # temperature of surrounding medium 

N = 50             # spatial discretization
a = R / N          # lattice constant
t, tmax = 0., 301. # time interval
h = .1             # time increment

# Construct an array to store the solution
T = Ti * np.ones(N + 3)  # N intervals => N+1 points + 2 ghost points


# Construct the matrices A and B and the vector C
alpha = D * h / (2 * a**2)
A = (1 + 2*alpha) * np.identity(N + 3)  # main diagonal of A
for n in range(1, N+1):                 # secondary diagonals
    A[n+1, n] = -alpha * (1 - 1/n)
    A[n+1, n+2] = -alpha * (1 + 1/n)
A[0, 0] = 1.   # first line (ghost): central derivative = 0
A[0, 2] = -1.
A[1, 0] = -3 * alpha  # second line, r=0: Laplacian = 3 d^2/dr^2
A[1, 1] = 1 + 6 * alpha
A[1, 2] = -3 * alpha
A[-1, -1] = 1. # last line (ghost): derivative ~ Delta T on the surface
A[-1, -2] = 2 * a * c
A[-1, -3] = -1
B = (1 - 2*alpha) * np.identity(N + 3)  # main diagonal of B
for n in range(1, N+1):                 # secondary diagonals
    B[n+1, n] = alpha * (1 - 1/n)
    B[n+1, n+2] = alpha * (1 + 1/n)
B[0, 0] = 0.  # first line (ghost), does not involve T(t)
B[1, 0] = 3 * alpha   # second line, r=0: Laplacian = 3 d^2/dr^2
B[1, 1] = 1 - 6 * alpha
B[1, 2] = 3 * alpha
B[-1, -1] = 0. # last line (ghost), does not involve T(t)
C = np.zeros(N+3) 
C[-1] = 2 * a * c * Tw  # the only nonzero element is the last one

# main loop
nextplot, dtplot = 0., 30.  # plot the solution every 30 seconds
r = 100 * np.linspace(0, R, N+1) # the values of the radial coordinate

while t <= tmax:
    if t >= nextplot:      # do we want to plot here? if yes, do so:
        plt.plot(r,  T[1:-1] - 273., label=f't = {t:.0f} s')
        nextplot += dtplot # next plot is again 30 seconds away
    t += h
    T = la.solve(A, B @ T + C) # this line is where the hard work is done
    
print(T - 273.)
plt.plot(np.linspace(5,7,2), [20, 20], 'k')              
plt.xlabel('r [cm]')
plt.ylabel('T [C]')
plt.text(5.5, 25, 'eau')
plt.legend()
plt.show()