Final project class

# -*- coding: utf-8 -*-
"""
Created on Thu May 23 22:47:11 2019

@author: Rode
"""

import numpy as np
import scipy.linalg as spl
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d

class Pendulum:
    def __init__(self, stepsize=0.2, Y_start = np.array([np.pi/2, 0, 0]), obst = -np.pi/6):
        self.stepsize = stepsize
        self.Y_points = np.array([Y_start])
        self.interpolvec = np.zeros((2,3))
        self.obst = obst

    def __repr__(self):
        """
        Returns the angular position and speed of the pendulum, as well as current time.
        """
        return "Position: {}    Speed: {}       Time: {}".format(self.Y_points[-1][0], self.Y_points[-1][1], self.Y_points[-1][2])

    def Eulerguess(self):
        """
        Makes a guess for next position and speed, using the explicit Euler method.
        """
        Y_guess = np.array([self.Y_points[-1][0] + self.stepsize*self.Y_points[-1][1],
                            self.Y_points[-1][1] + -self.stepsize*9.82*np.sin(self.Y_points[-1][0])])
        return Y_guess

    def Newton(self):
        Y_guess = self.Eulerguess()
        p = Y_guess[0]
        s = Y_guess[1]
        Y_delta = np.array([0,0])
        count = 0
        while 1e-10 < np.amax(Y_delta) or count < 5:
            count = count + 1
            J = np.array([    [1,   -self.stepsize/2],
                          [-np.cos(p)*self.stepsize/2,   1]   ])
            U = np.array([p - self.Y_points[-1][0] - (self.stepsize/2)*(self.Y_points[-1][1] + s),
                          s - self.Y_points[-1][1] + (self.stepsize/2)*9.82*    (np.sin(self.Y_points[-1][0]) + np.sin(p) ) ])
            Y_delta = spl.solve(J,-U)
            p = p + Y_delta[0]
            s = s + Y_delta[1]
            t = self.Y_points[-1][2] + self.stepsize
        return np.array([p, s, t])

    def goSteps(self, steps=1):
        for i in range(steps):
            self.Y_points = np.append(self.Y_points, [self.Newton()], axis=0)
            if 3 < len(self.Y_points):
                self.interpolvec = self.interpol3()
            if self.Y_points[-1,0] < self.obst:
                roots = np.roots([self.interpolvec[0,0], self.interpolvec[0,1], self.interpolvec[0,2] - self.obst])
                rootindex = np.where((self.Y_points[-1,2] <= roots) & (roots <= self.Y_points[-2,2]))
                print(rootindex)
                if len(rootindex) == 1:
                    obsttime = roots[(rootindex)]
                    print( obsttime)
                elif len(rootindex) == 2:
                    obsttime = np.max(roots)
                return obsttime

        return "Last step: Position: {}     Speed: {}   Time: {}".format(self.Y_points[-1][0], self.Y_points[-1][1],
                                                                         self.Y_points[-1][2])

    def Lagrange(self,X,Y):

        L = np.zeros((3,3))
        n = 3
        for j in range(n):
            Xj = np.delete(X,j)
            C = [X[j] - i for i in Xj]
            c = np.prod(C)
            polymult = [[1,-i] for i in Xj]
            J =polymult[0]
            for k in range(n-2):
                J = np.convolve(J, polymult[k+1])
            L[j] = J*Y[j]/c
        Poly = np.sum(L, axis=0)
#        Polyrev = [Poly[len(Poly) - 1 - i] for i in range(len(Poly))]
        return(Poly)

    def interpol3(self):
        X = self.Y_points[-3:,0]
        Y = self.Y_points[-3:,1]
        Z = self.Y_points[-3:,2]
        Poly1 = self.Lagrange(Z,X)
        Poly2 = self.Lagrange(Z,Y)
        Polyvec = np.array([Poly1, Poly2])
        return Polyvec

    def plot(self, choice):
        """
        What do you want to plot? Choose a number.
        1. x = time, y = position
        2. x = time, y = speed
        3. x = position, y = speed
        4. x = time, y1 = position (red), y2 = speed (blue)
        5. x = position, y = speed, z = time - 3D
        """
        if choice == 1:
            plt.plot(self.Y_points[...,2], self.Y_points[...,0])
        elif choice == 2:
            plt.plot([self.Y_points[...,2]], [self.Y_points[...,1]])
        elif choice == 3:
            plt.plot(self.Y_points[...,0], 'r', self.Y_points[...,1], 'b')
        elif choice == 4:
            plt.plot(self.Y_points[...,2], self.Y_points[...,0], 'r')
            plt.plot(self.Y_points[...,2], self.Y_points[...,1], 'b')
        elif choice == 5:
            plt.axes(projection='3d').plot3D(self.Y_points[...,0], self.Y_points[...,1], self.Y_points[...,2])
        plt.show()
        return