splines.py

import src.lagrange as lg
from sympy import symbols,diff,Matrix,lambdify,latex
import numpy as np
from numpy import linspace,array
import matplotlib.pyplot as plt
from src.dbg import dbg
from pandas import DataFrame

class SplineParab():
    def __init__(self):
        self.x=symbols('x')
        self.points=[]
        self.y=[]
        self.polinoms={}
        self.Lagrange=lg.Lagrange()
        self.n=0
        self.texPolinoms={}
        self.dataFrames=[]
    
    def showInNotebook(self):
        self.dataFrames=[]
        for i in range(self.n):
            #print(self.texPolinoms[(i,0)])
            self.dataFrames.append(DataFrame([self.texPolinoms[(i,0)]],columns=self.texPolinoms[(i,1)]))
            #self.plot(i)

    def plot(self,i=0):
        '''
        i for plot 
           p[i]   in points x[i],x[i+1]
           g      in points x[i+1],(x[i+1]+x[i+2])/2.
           h      in points (x[i+1]+x[i+2])/2.,x[i+2]
           p[i+1] in points x[i+2],x[i+3] 
        '''
        a0,a1,a2,a3=self.polinoms[(i,1)]
        p0,g,h,p1=self.polinoms[(i,0)]
        a=array(self.points)
        y=array(self.y)
        plt.plot(a,y,a0,p0(a0),a1,g(a1),a2,h(a2),a3,p1(a3))
        plt.legend(('p0(x)','g(x)','h(x)','p1(x)'),loc='upper left')
        plt.show()
        
    def splineParab(self,points=[0,1,2,3,4],y=[0,0,1,0,0],f='sin(x)'):
        '''
        DESCRIPTION:
        this function interpolate y use parabol splines
        f for case of ...
        
        
        TESTs:
        d=sp.splineParab(points=[0,1,2,3,4],y=[0,0,1,0,0])
        a0,a1,a2,a3=d[(0,1)]
        p0,g,h,p1=d[(0,0)]
        plt.plot(a0,p0(a0),a1,g(a1),a2,h(a2),a3,p1(a3))
        '''
        self.points=points
        self.y=y
        x=self.x
        self.polinoms={}
        self.texPolinoms={}
        self.n=len(points[:-3])
        for i in range(self.n):
            localPolinoms=[]
            localPoints=[]
            p0=self.Lagrange.lagrangeInterp(points=[points[i],points[i+1],points[i+2]],func=f,y=[y[i],y[i+1],y[i+2]])
            p1=self.Lagrange.lagrangeInterp(points=[points[i+1],points[i+2],points[i+3]],func=f,y=[y[i+1],y[i+2],y[i+3]])
            localPolinoms.append(p0)
            
            xHalf=(points[i+1]+points[i+2])/2.
            localPolinoms.extend(self.findParabs([points[i+1],xHalf,points[i+2]],[p0,p1]))
            localPolinoms.append(p1)
            
            localPoints.append(linspace(points[i],points[i+1],10))
            localPoints.append(linspace(points[i+1],xHalf,10))
            localPoints.append(linspace(xHalf,points[i+2],10))
            localPoints.append(linspace(points[i+2],points[i+3],10))
            
            self.polinoms[(i,0)]=[lambdify(x,p,"numpy") for p in localPolinoms]
            self.polinoms[(i,1)]=localPoints
            self.texPolinoms[(i,0)]=["$"+latex(p)+"$" for p in localPolinoms]
            self.texPolinoms[(i,1)]=[str([points[i],points[i+1]]),str([points[i+1],xHalf]),str([xHalf,points[i+2]]),str([points[i+2],points[i+3]])]

        #print(polinoms)
        return(self.polinoms)
        
    def findParabs(self,listX,listP):
        '''
        DESCRIPTION:
        find parabols equations g,h in points listX
        from that restrictions:
        g(x0)=p0(x0)
        g'(x0)=p0'(x0)
        h(x1)=p1(x1)
        h'(x1)=p1'(x1)
        g(xHalf)=h(xHalf)
        g'(xHalf)=h'(xHalf)
        '''
        x=self.x
        x0=listX[0]
        xHalf=listX[1]
        x1=listX[2]
        p0=listP[0]
        p1=listP[1]
        
        #x=symbols('x')
        a0=float(p0.subs(x,x0))
        a1=float(p1.subs(x,x1))
        b0=float((diff(p0,x)).subs(x,x0))
        b1=float((diff(p1,x)).subs(x,x1))
        dbg(**{"findParabs|part1":(a0,a1,b0,b1)})
        C=Matrix([[(xHalf-x0)**2,-(xHalf-x1)**2],[(xHalf-x0),-(xHalf-x1)]])
        #print("c=%s"%(C))
        b=Matrix([a1-a0+b1*(xHalf-x1)-b0*(xHalf-x0),(b1-b0)/2.])
        
        invC=C**-1
        
        c0,c1=invC*b
        dbg(**{"findParabs|part2":(b,invC,p0,p1,c0,c1)})
        
        g=a0+b0*(x-x0)+c0*(x-x0)**2
        h=a1+b1*(x-x1)+c1*(x-x1)**2
        return([g,h])