#!/usr/bin/env python
# coding: utf-8

# # Spring Optimization
# 
# 
# From: Sandgren, E. (1990). Nonlinear integer and discrete programming in mechanical design optimization. Transactions
# of the ASME, Journal of Mechanical Design, 112(2):223–229, June 1990. ISSN 0738-0666
# 
# **Global Optimum** (according to Lampinen): 2.65856
# 
# Note: In the runs below, you should be able to yield the exact global minimum, but in general, because a stochastic optimization method is being used, you will usually get close to, but not exactly to, the global minimum.  You may want to play with the control parameters of the optimization.

# In[1]:


import numpy as np
from math import pi
import PyCEGO


# In[2]:


# Problem originally from Sandgren
allowable_wire_diameters= [
    0.009, 0.0095, 0.0104, 0.0118, 0.0128, 0.0132,
    0.014, 0.015, 0.0162, 0.0173, 0.018, 0.020,
    0.023, 0.025, 0.028, 0.032, 0.035, 0.041,
    0.047, 0.054, 0.063, 0.072, 0.080, 0.092,
    0.105, 0.120, 0.135, 0.148, 0.162, 0.177,
    0.192, 0.207, 0.225, 0.244, 0.263, 0.283,
    0.307, 0.331, 0.362, 0.394, 0.4375, 0.500
]

# Constants of the problem
F_max = 1000.0 # [lb] Maximum working load
S = 189000.0 # [psi] Allowable maximum shear stress
l_max = 14.0 # [in] Maximum free length
d_min = 0.2 # [in] Minimum wire diameter
D_max = 3.0 # [in] Maximum outer diameter of the spring
F_p = 300.0 # [lb] Preload compression force
sigma_pm = 6.0 # [in] Maximum deflection under preload
sigma_w = 1.25 # [in] Deflection from preload to full load
G = 11.5e6 # [] Shear modulus of the material

def CEGO_obj(x):
    return obj([x[0].as_int(),
                x[1].as_double(),
                allowable_wire_diameters[x[2].as_int()]])

def obj(x):
    # Unpack the inputs
    N,D,d = x
    
    # Intermediate terms
    C_f = (4*(D/d)-1)/(4*(D/d)-4)+0.615*d/D
    K = G*d**4/(8*N*D**3)
    sigma_p = F_p/K
    l_f = F_max/K + 1.05*(N+2)*d
    
    g = np.zeros((9,))
    g[1] = 8*C_f*F_max*D/(pi*d**3) - S
    g[2] = l_f - l_max 
    # g3, g4, and g5 are handled explicitly via bounds; will always be met
    g[3] = d_min - d # will always be met
    g[4] = D-D_max # will always be met
    g[5] = 3.0-D/d # will always be met
    g[6] = sigma_p-sigma_pm
    g[7] = sigma_p+(F_max-F_p)/K + 1.05*(N+2)*d - l_f
    g[8] = sigma_w - (F_max - F_p)/K
    
    s = np.ones_like(g)
    s[7] = 1e5
    s[8] = 1e5
    
    c = np.array([1 if g_i < 0 else 1+s_i*g_i for s_i,g_i in zip(s,g)])
    assert((c >= 1).all())
    
    f = pi**2*D*d**2*(N+2)/4

    return f*np.product(c**3)

print('Quasi-optimal Values of the objective function according to the literature')
print('Sandgren:',obj([10, 1.180701, 0.283]), "(typo seems likely in Lampinen)") # Sandgren (see Lampinen)
print('Chen:',obj([9, 1.2287, 0.283])) # Chen and Tsao (see Lampinen)
print('Wu and Chao:',obj([9, 1.227411, 0.283])) # Wu and Chao (see Lampinen)
print('Lampinen:',obj([9, 1.2230410, 0.283])) # Lampinen


# In[3]:


D = 3
Nwired = len(allowable_wire_diameters)
CEGO_bounds = [PyCEGO.Bound(1, int(l_max/d_min)), # N
               PyCEGO.Bound(3*d_min, D_max), # D
               PyCEGO.Bound(27, Nwired-1) # d (indices thereof)
               ]

for ocounter in range(5):
    layers = PyCEGO.NumberishLayers(CEGO_obj, D, D*20, 1, 3)
    layers.set_bounds(CEGO_bounds)
    layers.set_builtin_evolver(PyCEGO.BuiltinEvolvers.differential_evolution)

    objs = []
    for counter in range(1000):
        layers.do_generation()
        objective, coeffs = layers.get_best()
        if counter % 50 == 0:
            print(layers.print_diagnostics())
        objs.append(objective)