#MSc Python
#Tutorial 1 - Debugging Functions & Things
## Department of Computing - Imperial College London ##
## Author: William Hunter ##

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import colors


###
###
### HOT ### 2-D SPACE + TIME EVOLUTION
###
###
###

### Things are heating up -  can you handle it?

#is this too low? 
resolution = 10**1

#same as medium
x_1 = np.linspace(1, 10, resolution)
x_2 = np.linspace(5, 50, resolution)

#same as medium
x_matrix = (x_1, x_2) #setting up 2-D vector space

#same as medium
print (x_matrix) #checking

#our function
def y_2Dfunc(x_1, x_2):
    y = np.sin(x_1) * np.cos(x_2) * x_1/x_2
    return

#is this working properly?
y_values = []
for i in range(len(x_1)):
    values = y_2Dfunc(x_1[i], x_2[i])
    y_values.append(values)

#let's plot against time now - treat it the same as x 
t_array = np.linspace(1, 100, resolution) #time array

#a new funciton that now calculates y in terms of time - you need to choose the right variables
def y_2Dfunc_time_evolution:
    y = (np.sin(x_1) * np.cos(x_2) * x_1/x_2) * np.sqrt(t_array) * np.log(t_array + 1)
    return y

#our for loop to generate the values by passing the right vectors through the funciton 
y_time = []
for i in range(len(t_array)):
    vals = y_2Dfunc_time_evolution(x_1, x_2, t_array[i])
    y_time.append(vals)


    
indep_range = np.linspace(1, 10, resolution)

#Here we are picking slices of our y_values for specific times  - defining our colour vectors 
#we need to do this as we will plot with a for-loop, as we have many slices!
picked_times = [y_time[1], y_time[2], y_time[5], y_time[10], y_time[20]]
times = (1, 2, 5, 10, 200)
colours = ('red', 'blue', 'green', 'pink', 'black', 'navy')
indep_range = np.linspace(1, 10, resolution)

#plotting the values
#here we are a number of slices and I want them plotted all on the same graph - can you do this with for loops??
#we're definitely on the spicy end of the sauces now...!
plt.figure(dpi = 100)
plt.xlabel('t')
plt.ylabel('y')
for i in range(len(picked_times)):
    plt.plot(indep_range, picked_times, '-', color = colours , label= 't = ' + (times[i]))
plt.legend(loc = 'best')#legend location = best
plt.show()