import numpy as np
from gekko import GEKKO

# Optimize at mesh points
x = np.arange(20.0, 30.0, 2.0)
y = np.arange(30.0, 50.0, 4.0)
amg, bmg = np.meshgrid(x, y)

# Initialize results array
obj = np.empty_like(amg)

m = GEKKO(remote=False)
objective = float('NaN')

a,b = m.Array(m.FV,2)

# model variables, equations, objective
x1 = m.Var(1,lb=1,ub=5)
x2 = m.Var(5,lb=1,ub=5)
x3 = m.Var(5,lb=1,ub=5)
x4 = m.Var(1,lb=1,ub=5)
m.Equation(x1*x2*x3*x4>=a)
m.Equation(x1**2+x2**2+x3**2+x4**2==b)
m.Minimize(x1*x4*(x1+x2+x3)+x3)
m.options.SOLVER = 1 # APOPT solver

# Calculate obj at all meshgrid points
for i in range(amg.shape[0]):
    for j in range(bmg.shape[1]):
        a.MEAS = amg[i,j]
        b.MEAS = bmg[i,j]

        m.solve(disp=False)

        obj[i,j] = m.options.OBJFCNVAL
        print(amg[i,j],bmg[i,j],obj[i,j])

# plot 3D figure of results
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(amg, bmg, obj, \
       rstride=1, cstride=1, cmap=cm.coolwarm, \
       vmin = 10, vmax = 25, linewidth=0, antialiased=False)
ax.set_xlabel('a')
ax.set_ylabel('b')
ax.set_zlabel('obj')
plt.show()