# -*- coding: utf-8 -*-
"""
BYU Intro to Optimization. Column design
https://apmonitor.com/me575/index.php/Main/TubularColumn
Contour plot additions by Al Duke 11/30/2019
"""
from gekko import GEKKO
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import fsolve

m = GEKKO()

#%% Constants
pi = m.Const(3.14159,'pi')
P = 2300 # compressive load (kg_f)
o_y = 450 # allowable yield stress (kg_f/cm^2)
E = 0.65e6 # elasticity (kg_f/cm^2)
p = 0.0020 # weight density (kg_f/cm^3)
l = 300 # length of the column (cm)

#%% Variables (the design variables available to the solver)
d = m.Var(value=8.0,lb=2.0,ub=14.0) # mean diameter (cm)
t = m.Var(value=0.3,lb=0.2 ,ub=0.8) # thickness (cm)
cost = m.Var()

#%% Intermediates (computed by solver from design variables and constants)
d_i = m.Intermediate(d - t)
d_o = m.Intermediate(d + t)
W = m.Intermediate(p*l*pi*(d_o**2 - d_i**2)/4) # weight (kgf)
o_i = m.Intermediate(P/(pi*d*t)) # induced stress

# second moment of area of the cross section of the column
I = m.Intermediate((pi/64)*(d_o**4 - d_i**4))

# buckling stress (Euler buckling load/cross-sectional area)
o_b = m.Intermediate((pi**2*E*I/l**2)*(1/(pi*d*t)))

#%% Equations (constraints, etc. Cost could be an intermediate variable)
m.Equations([
o_i - o_y <= 0,
o_i - o_b <= 0,
cost == 5*W + 2*d
])

#%% Objective
m.Minimize(cost)

#%% Solve and print solution
m.options.SOLVER = 1
m.solve()

print('Optimal cost: ' + str(cost[0]))
print('Optimal mean diameter: ' + str(d[0]))
print('Optimal thickness: ' + str(t[0]))

minima = np.array([d[0], t[0]]) 

#%% Contour plot
# create a cost function as a function of the design variables d and t
f = lambda d, t: 2 * d + 5 * p * l * np.pi * ((d+t)**2 - (d-t)**2)/4 

xmin, xmax, xstep = 2, 14, .2 # diameter
ymin, ymax, ystep = .2, .8, .05 # thickness

d, t = np.meshgrid(np.arange(xmin, xmax + xstep, xstep), \
                   np.arange(ymin, ymax + ystep, ystep))
z = f(d, t)

# Determine the compressive stress constraint line. 
#stress = P/(pi*d*t) # induced axial stress
t_stress = np.arange(ymin, ymax, .025) # use finer step to get smoother constraint line
d_stress = []
for tt in t_stress:
    dd = P/(np.pi * tt * o_y)
    d_stress.append(dd)

# Determine buckling constraint line. This is tougher because we cannot
#  solve directly for t from d. Used scipy.optimize.fsolve to find roots 
d_buck = []
t_buck = []
for d3 in np.arange(6, xmax, .005): 
    fb = lambda t : o_y-np.pi**2*E*((d3+t)**4-(d3-t)**4)/(64*l**2*d3*t)
    tr = np.array([0.3])
    roots = fsolve(fb, tr)
    if roots[0] != 0: 
        if roots[0] >= .1 and roots[0]<=1.:
            t_buck.append(roots[0])
            d_buck.append(d3)

# Create contour plot
plt.style.use('ggplot') # to make prettier plots
fig, ax = plt.subplots(figsize=(10, 6))

CS = ax.contour(d, t, z, levels=15,) 
ax.clabel(CS, inline=1, fontsize=10)
ax.set_xlabel('mean diameter $d$')
ax.set_ylabel('half thickness $t$')

ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))

# Add constraint lines and optimal marker
ax.plot(d_stress, t_stress, "->", label="Stress constraint")
ax.plot(d_buck, t_buck, "->", label="Buckling constraint" )

minima_ = minima.reshape(-1, 1)
ax.plot(*minima_, 'r*', markersize=18, label="Optimum")
ax.text(10,.25,"Contours = Cost (objective)\nConstraint line markers point\ntowards feasible space.")
plt.title('Column Design')
plt.legend()
plt.show()