Transfer Functions in python

INTERACTIVE_PLOT = True
if INTERACTIVE_PLOT:
    %matplotlib notebook
    pass
else:
    %matplotlib inline
    pass

import sympy
from sympy import *
sympy.init_printing()
s = Symbol('s')
z = Symbol('z')

from control import matlab
from control import pzmap

import matplotlib.pyplot as plt

DEFAULT_DT = 0.001
#Converts a polynomial in z to a transfer function
def tfDiscrete(Ts, dt = DEFAULT_DT, *args, **kwargs):
    tfunc = Ts.simplify() #This is necessary, otherwise you can get float errors (the result is too inexact.)
    num = Poly(tfunc.as_numer_denom()[0],z).all_coeffs()
    den = Poly(tfunc.as_numer_denom()[1],z).all_coeffs()
    tf = matlab.tf(map(float,num),map(float,den), dt)
    return tf

def tfCont(Ts, *args, **kwargs):
    tfunc = Ts.simplify() #This is necessary, otherwise you can get float errors (the result is too inexact.)
    num = Poly(tfunc.as_numer_denom()[0],s).all_coeffs()
    den = Poly(tfunc.as_numer_denom()[1],s).all_coeffs()
    tf = matlab.tf(map(float,num),map(float,den))
    return tf

def tf(Ts, dt = DEFAULT_DT, *args, **kwargs):
    if len(Ts.free_symbols) > 1:
        raise ValueError('Too many free variables, ' + str(Ts.free_symbols) +
                         ' A transfer function is a polynomial in only s or z.')
    if len(Ts.free_symbols) < 1:
        raise ValueError('Too few variables.' 
                         'A transfer function is a polynomial in s or z.')
    if z in Ts.free_symbols :
        return tfDiscrete(Ts, dt, *args, **kwargs)
    elif s in Ts.free_symbols:
        return tfCont(Ts, *args, **kwargs)
    else:
        raise ValueError('A transfer function is a polynomial in s or z.'
                        'not one in ' + str(Ts.free_symbols))
def pole(Ts, dt = DEFAULT_DT, *args,**kwargs):
    return matlab.pole(tf(Ts,dt),*args,**kwargs)

def rlocus(Ts, dt = DEFAULT_DT, *args,**kwargs):
    plt.figure()
    matlab.rlocus(tf(Ts,dt),*args,**kwargs)
    plt.show()

def bode(Ts, dt = DEFAULT_DT, *args,**kwargs):
    plt.figure()
    matlab.bode(tf(Ts,dt),*args,**kwargs)
    plt.show()

def polezero(Ts, dt = DEFAULT_DT):
    plt.figure()
    pz = pzmap.pzmap(tf(Ts,dt))
    plt.show()
    return pz

def damp(Ts, dt = DEFAULT_DT, *args,**kwargs):
    return matlab.damp(tf(Ts,dt),*args,**kwargs)

def stepResponse(Ts, dt = DEFAULT_DT, *args,**kwargs):
    plt.figure()
    tfunc = tf(Ts,dt,*args,**kwargs)
    y,t = matlab.step(tfunc,*args,**kwargs)
    if(len(t)==len(y)): # Continuous time
        plt.plot(t,y)
        plt.title("Step Response")
        plt.grid()
        plt.xlabel("time (s)")
        plt.ylabel("y(t)")
        info ="Over Shoot: %f%s"%(round((y.max()/y[-1]-1)*100,2),'%')
        try:
            i10 = next(i for i in range(0,len(y)-1) if y[i]>=y[-1]*.10)
            Tr = round(t[next(i for i in range(i10,len(y)-1) if y[i]>=y[-1]*.90)]-t[i10],2)
        except StopIteration:
            Tr = "unknown"
        try:
            Ts = round(t[next(len(y)-i for i in range(2,len(y)-1) if abs(y[-i]/y[-1])>1.02 or abs(y[-i]/y[-1])<0.98)]-t[0],2)
        except StopIteration:
            Ts = "unknown"

        info += "\nRise Time: %s"%(Tr)
        info +="\nSettling time: %s"%(Ts)
        print info
        plt.legend([info],loc=4)
        plt.show()
    else: #discrete time 
        y = y[0] #unpack value
        t = [x*dt for x in range(len(y))]
        plt.plot(t, y)
        plt.title("Step Response")
        plt.grid()
        plt.xlabel("time (s)")
        plt.ylabel("y[t]")
        info ="Over Shoot: %f%s"%(round((y.max()/y[-1]-1)*100,2),'%')
        try:
            i10 = next(i for i in range(0,len(y)-1) if y[i]>=y[-1]*.10)
            Tr = round(t[next(i for i in range(i10,len(y)-1) if y[i]>=y[-1]*.90)]-t[i10],2)
        except StopIteration:
            Tr = "unknown"
        try:
            Ts = round(t[next(len(y)-i for i in range(2,len(y)-1) if abs(y[-i]/y[-1])>1.02 or abs(y[-i]/y[-1])<0.98)]-t[0],2)

        except StopIteration:
            Ts = "unknown"

        info += "\nRise Time: %s"%(Tr)
        info +="\nSettling time: %s"%(Ts)
        plt.legend([info],loc=4)
        plt.show()
def impulseResponse(Ts, dt = DEFAULT_DT, *args,**kwargs):
    plt.figure()
    tfunc = tf(Ts,dt,*args,**kwargs)
    y,t = matlab.impulse(tfunc,*args,**kwargs)
    if(len(t)==len(y)): # Continuous time
        plt.plot(t,y)
        plt.title("Impulse Response")
        plt.grid()
        plt.xlabel("time (s)")
        plt.ylabel("y(t)")
        plt.show()
    else: #discrete time 
        y = y[0] #unpack value
        t = [x*dt for x in range(len(y))]
        plt.plot(t, y)
        plt.title("Impulse Response")
        plt.grid()
        plt.xlabel("time (s)")
        plt.ylabel("y[t]")
        #info ="Over Shoot: %f%s"%(round((y.max()/y[-1]-1)*100,2),'%')
        #plt.legend([info],loc=4)
        plt.show()

def oscillationPeriod(Ts, dt = DEFAULT_DT, *args,**kwargs):
    a = pole(Ts,0.001)
    return 2*3.1415/max(a, key=lambda x: abs(x)).imag

Kp = Symbol('Kp')
Kd = Symbol('Kd')
Ki = Symbol('Ki')
p = Symbol('p')
P = Symbol('P')
m = Symbol('m')
gamma = Symbol('y')
dt = Symbol('dt')

ha2w = dt/(z-1)
hc2a = gamma
hw2th = dt/(z-1)
hc2ap = (1-p) *(z**-1)/(1-p*z**-1)*gamma

H = hc2a*ha2w*hw2th
Hp = hc2ap*ha2w*hw2th

actuator = H
actuatorP = Hp

controller = (Kp+Kd/(dt*m)-Kd/(dt*m)*z**(-m))
sensor = P

G=controller*actuator; 
Gp=controller*actuatorP; 

sys = G/(1+G*sensor)
sysP = Gp/(1+Gp*sensor)

sys = sys.simplify()
sysP = sysP.simplify()

pprint(sys)
print("\n\n ------------------------- \n\n")
pprint(sysP)

#Calculating the value of p
# N number of steps untill we are 50% done. 1 - p**n = 1/2. p = 1/2.*(1./n)
n = 65.
pval = 1/2.**(1/n)
print pval 

#We want the period to be 530 when Kp=10 and Kd=1
#By trial and error:
subs = {gamma:13.4, 
        Kp:10, 
        Kd:1, 
        Ki:0, 
        m:3, 
        P:1, 
        p:pval,  
        dt:0.001}

print(oscillationPeriod(sysP.subs(subs),0.001))

# Now we want to find the best Kp and Kd
# Brute Force approach
a=[]
for kp in range(1,40):
    for kd in range(1,40):
        kpp = kp/2.
        kdd = kd/10.
        subs[Kp] = kpp
        subs[Kd] = kdd
        a.append((kpp,kdd,max(abs(pole(sysP.subs(subs),0.001)))))
kpp, kdd , magnitude = min(a, key=lambda x: x[2])
print (kpp, kdd , magnitude)

# Non Bruteforce
from scipy.optimize import minimize

def func_min(x):
    kpp = x[0]
    kdd = x[1]
    subs[Kp] = kpp
    subs[Kd] = kdd
    return max(abs(pole(sysP.subs(subs),0.001)))

x0 = [8, 1]
res = minimize(func_min, x0, method='nelder-mead', options={'disp': True})
print res

# Comparing both results
print subs
subs[Kp] = 1.5
subs[Kd] = 0.6
stepResponse(sysP.subs(subs), 0.001, T=[x*0.001 for x in range(3000)])
impulseResponse(sysP.subs(subs), 0.001, T=[x*0.001 for x in range(3000)])
bode(sysP.subs(subs),0.001)
print polezero(sysP.subs(subs),0.001)


subs[Kp] = 0.30533799
subs[Kd] = 0.26041388
stepResponse(sysP.subs(subs), 0.001, T=[x*0.001 for x in range(3000)])
impulseResponse(sysP.subs(subs), 0.001, T=[x*0.001 for x in range(3000)])
bode(sysP.subs(subs),0.001)
print polezero(sysP.subs(subs),0.001)