# -*- coding: utf-8 -*-
"""
Created on Sun Oct  4 14:44:38 2020

@author: jbobowsk
"""

# In this Python script we attempt to solve a partial differential equation (PDE).
# We will use the 'FiPy' module to do this.  I found that the module was quite
# difficult to learn.  A lot of the code below is based on code shown in the
# FiPy documentation.  I made a few tweaks here and there.

# Note that FiPy is note one of the standard modules and you will likely
# have to install it manually in order to use it.  I'm using the Anaconda
# platform.  To install the FiPy module, I opended the "Anaconda Powershell
# Prompt" and entered "conda install - c conda-forge fipy" (without the 
# quotes).
from fipy import *
import matplotlib.pyplot as plt
import numpy as np

# We will attempt to solve the 1-D heat equation, also  known as the 1-D 
# diffusion equation:
# alpha*T" = dT/dt
# Here, T is temperature, alpha is thermal diffusivity (a diffusion coefficient),
# and t is time.  T" represents two spatial derivative of temperature.

# We will solve for temperature profiles along a 1-D rod subject to particular
# boundary conditions and some initial temperature profile.  We will get solutions
# at a bunch of different times between t = 0 (our initial condition) and a final
# time.  We need to set up a spatial grid that defines the length of our rod and
# the number of points between the two ends for which we will calculate the
# temperature.  We will have a rod that is 1.5 m long and we will extract the 
# temperature at 50 equally-spaced points.
nx = 50
dx = 1.5/nx
mesh = Grid1D(nx = nx, dx = dx)

# Let's set up a list of x-coordinates to be used later for plotting purposes.
xx = np.arange(0, nx*dx, dx)
print(xx)

# Define the initial temperature profile of the bar.  We will assume that the
# bar is initially at a uniform temperature of 20 celsius.
temperature = CellVariable(name = '', mesh = mesh, value = 20.)

# The thermal diffusivity of copper near room temperature is about
# 0.0001 m^2/s
alpha = 1e-4 
print('The thermal diffusivity of copper near romm temperature is', alpha, 'm^2/s.')

# These are the boundary conditions, i.e. the temperatures in celsius of the 
# two ends of the bar for t > 0.
leftEnd = 100
rightEnd = 20

# Here's the syntax that FiPy uses to define the boundary conditions.
temperature.constrain(rightEnd, mesh.facesRight)
temperature.constrain(leftEnd, mesh.facesLeft)

# Here's how we define the partial differential equation (PDE) that we hope to
# solve.  'TransientTerm()' represents the dT/dt and 'ExplicitDiffusionTerm()'
# represents the T" with a coefficient of alpha.
eqX = TransientTerm() == ExplicitDiffusionTerm(coeff = alpha)

# The 'timeStepDuration' that has been selected was just copied from the 
# example in the FiPy documentation.  It represents the amount of time
# between successive solutions for the temperature profile.  We can run the 
# solver for as many steps as we want.  I selected 2500 steps which gives
# the bar enough time to reach the expected equilibrium temperature profile. 
timeStepDuration = 0.9*dx**2/(2*alpha)
print('The time between steps is', timeStepDuration, 'seconds.')
steps = int(2.5e3)
print('The total elapsed time is', steps*timeStepDuration/3600, 'hours.')

# A line a code that I had to look up online.  The 'viewer' that has been defined
# will allow use to look that the intermediate solutions for the temperature
# profiles as the loop below advances the solutions from the intial time all the
# way through to step 2500.
viewer = MatplotlibViewer(vars=(temperature,), datamin=15, datamax=110, legend=None)

# Here's how we call the 'viewer'.  This plot will show use the temperature
# profile at t = 0.  It will just be a uniform temperatire of 20 C.
print('A plot of the initial temperature profile of the copper bar.')
viewer.plot()

# I couldn't figure out how to extract the actual temperature data from the FiPy
# solver at the individual solution times.  After a lot of time on Google, I came
# up with a work around.  The 'TSVViewer' function will allow us to write the 
# desired data to a file.  I will then read that file back into Python so that
# it can be stored in an array.  I assume that there must be a better way...
from fipy.viewers.tsvViewer import TSVViewer  

# Here's the for loop that will step through the different times and then call
# the PDE solver at each iteration.  'solnData[]' will be a list of the temperature
# profiles to be used for plotting later.
solnData = []
for step in range(steps):
     # Here's where we call the PDE solver.
     eqX.solve(var = temperature, dt = timeStepDuration)
     # I don't want to see the temperature profiles at all 2500 of the solution
     # times.  This if statement selects 20 solutions, uniformly spaced in time
     # to look at.
     if step % int(steps/20) == 0:
         # The expected equilibrium solution is a linear temperature profile
         # starting from 100 C at the left end and ending at 20 C at the right
         # end.  We will add this line to out plots so that we can see how our
         # solutions approach equilibium as time evolves.
         plt.plot(xx, leftEnd - (leftEnd - rightEnd)/(nx*dx)*xx, 'k--', linewidth = 0.5)
         # Next, we call 'viewer' to displat the relevant temperature profile.
         viewer.plot()
         print('The time is', step*timeStepDuration/3600, 'hours.')
         # Then, TSVViewer is used the write the data to a file.  This file will
         # be overwritten each time the if condition is true.  If you wanted to
         # keep the data, you could modify the code so that the file name is
         # changed each time 'TSVViewer' is called.
         TSVViewer(vars=(temperature)).plot('data1.txt')
         # This next block of code opens the file that was just created and then
         # reads it into memory line by line.  We then load the data into 'data'
         # using 'np.loadtex()' while skipping the first line.  The first line
         # needs to be skipped because TSVViewer adds a one-line header to the file.
         # When the indentation after the 'with' statement ends, the file is closed.
         with open('data1.txt') as f:
             lines = (line for line in f if not line.startswith('#'))
             data = np.loadtxt(lines, skiprows = 1)[:, 1]
         # Add the data that we just extracted in a convoluted way to a list
         # of solutions.
         solnData = solnData + [data]
     elif step == steps - 1:
         # Do all of those same steps if we've reached that last step in the loop.
         plt.plot(xx, leftEnd - (leftEnd - rightEnd)/(nx*dx)*xx, 'k--', linewidth = 0.5)
         viewer.plot()
         print('The solution at the final time of', step*timeStepDuration/3600, 'hours.')
         with open('data1.txt') as f:
             lines = (line for line in f if not line.startswith('#'))
             data = np.loadtxt(lines, skiprows = 1)[:, 1]
         solnData = solnData + [data]

# Finally, we can now plot all of the temperature profiles at the different
# times on a single graph.             
plt.figure(figsize=(15,15))

# We'll use a colormap (cm) to nicely color the successive curves.
from matplotlib.pyplot import cm
color = cm.hot(np.linspace(0,1,len(solnData)))
for i in range(len(solnData)):
    plt.plot(xx, solnData[i], c = color[i], linewidth = 2, alpha = 0.75)   
plt.xlabel('position along the copper bar (m)')
plt.ylabel('temperature (c)')

# We didn't have to start with a rod at a uniform temperature at t = 0.
# We can start with whatever initial temperature profile we want.  
# Here's a linear/sinusoidal initial temperature profile.
temperature = CellVariable(name='', mesh=mesh, value=(40*xx - 50*np.sin(2*np.pi*xx/(nx*dx))))

# Besides, this change, we'll run all of the same code from above again, this
# time without all of the comments.
viewer = MatplotlibViewer(vars=(temperature,), datamin=-60, datamax=110, legend=None)
print('A plot of the initial temperature profile of the copper bar.')
viewer.plot()

solnData = []
for step in range(steps):
     eqX.solve(var = temperature, dt = timeStepDuration)
     if step % int(steps/20) == 0:
         plt.plot(xx, leftEnd - (leftEnd - rightEnd)/(nx*dx)*xx, 'k--', linewidth = 0.5)
         viewer.plot()
         print('The time is', step*timeStepDuration/3600, 'hours.')
         TSVViewer(vars=(temperature)).plot('data1.txt')
         with open('data1.txt') as f:
             lines = (line for line in f if not line.startswith('#'))
             data = np.loadtxt(lines, skiprows = 1)[:, 1]
         solnData = solnData + [data]
     elif step == steps - 1:
         plt.plot(xx, leftEnd - (leftEnd - rightEnd)/(nx*dx)*xx, 'k--', linewidth = 0.5)
         viewer.plot()
         print('The solution at the final time of', step*timeStepDuration/3600, 'hours.')
         with open('data1.txt') as f:
             lines = (line for line in f if not line.startswith('#'))
             data = np.loadtxt(lines, skiprows = 1)[:, 1]
         solnData = solnData + [data]
plt.figure(figsize=(15,15))
color = cm.hot(np.linspace(0,1,len(solnData)))
for i in range(len(solnData)):
    plt.plot(xx, solnData[i], c = color[i], linewidth = 2, alpha = 0.75)   
plt.xlabel('position along the copper bar (m)')
plt.ylabel('temperature (c)')