import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
from gekko import GEKKO
from matplotlib.widgets import Button
import time
tf = 200
t_span = np.linspace(0, tf, tf + 1)
valve = 0.0
# Setpoint trajectory
sp = np.ones(tf + 1) *0.5
sp[int(tf/2):] = np.random.choice([0.2, 0.3, 0.4, 0.6, 0.7, 0.8])
# Create plots
scale = 0.5
fig, axs = plt.subplots(3, 2, figsize=(16*scale, 9*scale), sharey='row', gridspec_kw={'width_ratios': [1, 2]})
# Initialize bar plots
tank_1_bar, = axs[0, 0].bar(x = 0, height = 0, width = 1.0, color = 'b', label='Tank 1 Level')
tank_2_bar, = axs[1, 0].bar(x = 0, height = 0, width = 1.0, color = 'b', label='Tank 2 Level')
pump_speed, = axs[2, 0].bar(x = 0, height = 0.5, width = 1.0, facecolor = 'none', edgecolor = 'black', label='Pump speed')
# Initialize line plots
line_y1, = axs[0, 1].plot([], [], 'b-', label='Manual')
line_y1_gekko, = axs[0, 1].plot([], [], 'b--', label='GEKKO MPC')
line_y2, = axs[1, 1].plot([], [], 'r-', label='Manual')
line_y2_gekko, = axs[1, 1].plot([], [], 'r--', label='GEKKO MPC')
line_sp, = axs[1, 1].plot([], [], 'k-', label='Setpoint')
line_pump, = axs[2, 1].plot([], [], 'g-', label='Manual')
line_pump_gekko, = axs[2, 1].plot([], [], 'g--', label='GEKKO MPC')
# Set axes ranges and labels
axs[0, 1].legend()
axs[1, 1].legend()
axs[2, 1].legend()
axs[0, 0].set_ylabel('Tank 1 level')
axs[0, 0].set_xlim(-0.5,0.5)
axs[0, 0].set_ylim(-0.1,1.1)
axs[0, 0].set_title = 'Tank 1 level'
axs[1, 0].set_ylabel('Tank 2 level')
axs[1, 0].set_xlim(-0.5,0.5)
axs[1, 0].set_ylim(-0.1,1.1)
axs[2, 0].set_ylabel('Pump speed')
axs[2, 0].set_xlim(-0.5,0.5)
axs[2, 0].set_ylim(-0.1,1.6)
axs[0, 1].set_xlim(0,tf)
axs[1, 1].set_xlim(0,tf)
axs[2, 1].set_xlim(0,tf)
axs[2, 1].set_xlabel('Time')
# Show the tank and pump limit lines
axs[0,0].plot([-0.5, 0.5], [1.0, 1.0], 'k:')
axs[1,0].plot([-0.5, 0.5], [1.0, 1.0], 'k:')
axs[2,0].plot([-0.5, 0.5], [1.5, 1.5], 'k:')
axs[0,1].plot([0, tf], [1.0, 1.0], 'k:')
axs[1,1].plot([0, tf], [1.0, 1.0], 'k:')
axs[2,1].plot([0, tf], [1.5, 1.5], 'k:')
# Create button axes
ax_increase = plt.axes([0.8, 0.02, 0.1, 0.04])
ax_decrease = plt.axes([0.7, 0.02, 0.1, 0.04])
# Create buttons
btn_increase = Button(ax_increase, 'Pump +0.05')
btn_decrease = Button(ax_decrease, 'Pump -0.05')
plt.ion()
plt.show()
# Simulated system (for measurements)
def tank(levels, t, pump, valve):
h1 = max(1.0e-10, levels[0])
h2 = max(1.0e-10, levels[1])
c1 = 0.08 # inlet valve coefficient
c2 = 0.04 # tank outlet coefficient
dhdt1 = c1 * (1.0 - valve) * pump - c2 * np.sqrt(h1)
dhdt2 = c1 * valve * pump + c2 * np.sqrt(h1) - c2 * np.sqrt(h2)
if h1 >= 1.0 and dhdt1 > 0.0:
dhdt1 = 0
if h2 >= 1.0 and dhdt2 > 0.0:
dhdt2 = 0
return [dhdt1, dhdt2]
def sim():
# Manual simulation loop
current_i = 0
# Initial conditions
h0 = [0, 0]
# Record the solution
y = np.zeros((tf + 1, 2))
y[0, :] = h0
# Initial pump values
current_pump = 0.5 # Initial pump value
pump = np.ones(tf + 1) * current_pump # Initialize pump array
# Button callback functions
def increase_pump(event):
nonlocal current_pump, pump, current_i
current_pump = min(current_pump + 0.05, 1.5)
pump[current_i:] = current_pump
def decrease_pump(event):
nonlocal current_pump, pump, current_i
current_pump = max(current_pump - 0.05, 0.0)
pump[current_i:] = current_pump
# Connect buttons
btn_increase.on_clicked(increase_pump)
btn_decrease.on_clicked(decrease_pump)
for i in range(tf):
current_i = i
# Integrate the model
inputs = (pump[i], valve)
h = odeint(tank, y[i, :], [0, 1], inputs)
# Record the result
y[i + 1, :] = h[-1, :]
# Update bar plot data
tank_1_bar.set_height(y[i + 1, 0])
tank_2_bar.set_height(y[i + 1, 1])
pump_speed.set_height(pump[i])
# Change the bar colors if the tanks overflow
if y[i + 1, 0] < 1.0:
tank_1_bar.set_color('b')
else:
tank_1_bar.set_color('r')
if y[i + 1, 1] < 1.0:
tank_2_bar.set_color('b')
else:
tank_2_bar.set_color('r')
# Update line data
line_sp.set_data(t_span[:i + 1], sp[:i + 1])
line_y1.set_data(t_span[:i + 1], y[:i + 1, 0])
line_y2.set_data(t_span[:i + 1], y[:i + 1, 1])
line_pump.set_data(t_span[:i + 1], pump[:i + 1])
plt.pause(0.015)
# Allow some time adjust the figure area
if i == 0:
plt.pause(5)
## GEKKO MPC Solution ##
# Reset initial conditions
h0 = [0,0]
y = np.zeros((tf+1,2))
y[0,:] = h0
pump = np.zeros(tf+1)
# create MPC with GEKKO
m = GEKKO(remote = False)
m.time = [0, 2, 4, 8, 10, 20]
# empirical constants
Kp_h1 = 1.4
tau_h1 = 18.4
Kp_h2 = 1
tau_h2 = 24.4
# manipulated variable (pump)
p = m.MV(value=0.5,lb=1e-5,ub=1.5)
p.STATUS = 1
p.DCOST = 0.01
p.FSTATUS = 0
# unmeasured state
h1 = m.Var(value=0.0, ub = 0.90)
# controlled variable
h2 = m.CV(value=0.0)
h2.STATUS = 1
h2.FSTATUS = 1
h2.TAU = 20
h2.TR_INIT = 1
# equations
m.Equation(tau_h1*h1.dt()==-h1 + Kp_h1*p)
m.Equation(tau_h2*h2.dt()==-h2 + Kp_h2*h1)
# options
m.options.IMODE = 6
m.options.CV_TYPE = 2
m.options.MV_STEP_HOR = 1
# GEKKO simulation loop
current_i = 0
for i in range(tf):
current_i = i
#########################
# MPC ###################
#########################
# measured height
h2.MEAS = y[i,1]
# set point deadband
h2.SPHI = sp[i]+0.01
h2.SPLO = sp[i]-0.01
h2.SP = sp[i]
# solve MPC
m.solve(disp=False)
# retrieve 1st pump new value
pump[i] = p.NEWVAL
#########################
# System ################
#########################
# Specify the pump and valve
inputs = (pump[i],valve)
# Integrate the model
h = odeint(tank,h0,[0,1],inputs)
# Record the result
y[i+1,:] = h[-1,:]
# Reset the initial condition
h0 = h[-1,:]
# Update bar plot data
tank_1_bar.set_height(y[i + 1, 0])
tank_2_bar.set_height(y[i + 1, 1])
# Change the bar colors if the tanks overflow
if y[i + 1, 0] < 1.0:
tank_1_bar.set_color('b')
else:
tank_1_bar.set_color('r')
if y[i + 1, 1] < 1.0:
tank_2_bar.set_color('b')
else:
tank_2_bar.set_color('r')
pump_speed.set_height(pump[i])
# Update line data
line_y1_gekko.set_data(t_span[:i], y[:i, 0])
line_y2_gekko.set_data(t_span[:i], y[:i, 1])
line_pump_gekko.set_data(t_span[:i], pump[:i])
plt.pause(0.01)
run_sim = True
while run_sim:
sim()
z = input('Run again (Y/N)?')
run_sim = True if z.upper() == 'Y' else False
if run_sim:
# Reset bar plots
tank_1_bar.set_height(0)
tank_2_bar.set_height(0)
pump_speed.set_height(0.5)
# Reset line plots
line_y1.set_data([], [])
line_y1_gekko.set_data([], [])
line_y2.set_data([], [])
line_y2_gekko.set_data([], [])
line_sp.set_data([], [])
line_pump.set_data([], [])
line_pump_gekko.set_data([], [])