1. import math
2. import numpy as np
3. from scipy.optimize import fmin
4. N = 6
5. x0 = 1
6.  
7. # Ranges:
8. # x in [0, ..., N)
9. # u in [0, ..., N - 1)
10. # v in [0, ..., N)
11.  
12.  
13. def get_state_values(u: np.ndarray) -> np.ndarray:
14.     x = np.zeros((N))
15.     x[0] = x0
16.     for i in range(1, N):
17.         x[i] = 2 * x[i - 1] + u[i - 1]
18.     return x
19.  
20.  
21. def constrained_f(u: np.ndarray, v: np.ndarray, t: float) -> float:
22.     x = get_state_values(u)
23.     f = 0
24.     for i in range(N - 1):
25.         f += math.pow(x[i], 2) + 2 * math.pow(u[i], 2)
26.     cost = 0
27.     # Constraint imposed on control value:
28.     for i in range(N - 1):
29.         cost += (u[i] - v[i, 0]) * max(0, u[i] - v[i, 0])
30.         cost += (-u[i] - v[i, 1]) * max(0, -u[i] - v[i, 1])
31.     # Constraint imposed on the final state value:
32.     cost += math.pow((x[N-1] - v[N - 1, 0]), 2)
33.     return f + (t/2) * cost
34.  
35.  
36. def main():
37.     alpha = 2
38.     beta = 2
39.     epsilon = 0.1
40.     t = 0.1
41.     c = 0.1
42.     x = np.zeros((N))
43.     u = np.zeros((N - 1))
44.     v = np.zeros((N, 2))
45.     lower = 1
46.     upper = 3
47.     last = 20
48.     for i in range(N - 1):
49.         v[i, 0] = upper
50.         v[i, 1] = lower
51.     v[N - 1][0] = last
52.  
53.     while True:
54.         f_old = constrained_f(u, v, t)
55.         u = fmin(constrained_f, u, args=(v, t), disp=False)
56.         f_new = constrained_f(u, v, t)
57.  
58.         r = np.zeros((N, 2))
59.         for i in range(N - 1):
60.             r[i, 0] = u[i] - upper
61.             r[i, 1] = -u[i] - lower
62.         r[N - 1, 0] = x[N - 1] - last
63.  
64.         gamma = np.zeros((N, 2))
65.         for i in range(N - 1):
66.             gamma[i, 0] = v[i, 0] + r[i, 0] - upper
67.             gamma[i, 1] = v[i, 1] + r[i, 1] - lower
68.         gamma[N - 1, 0] = x[N - 1] + r[N - 1, 0] - last
69.         gamma_sum = np.sum(gamma)
70.         if gamma_sum <= epsilon:
71.             if np.abs(r[N - 1, 0]) <= np.abs(gamma[N - 1, 0])
72.                     and np.abs(f_old - f_new) <= epsilon:
73.                 break
74.         else:
75.             if gamma_sum < c:
76.                 for i in range(N - 1):
77.                     v[i, 0] = upper - r[i, 0]
78.                     v[i, 1] = lower - r[i, 1]
79.                 v[N - 1, 0] = last - r[N - 1, 0]
80.                 c = alpha * gamma_sum
81.             else:
82.                 t = beta * t
83.                 for i in range(N - 1):
84.                     v[i, 0] = upper - r[i, 0] / beta
85.                     v[i, 1] = lower - r[i, 1] / beta
86.                 v[N - 1, 0] = last - r[N - 1, 0] / beta
87.  
88.     print(u)
89.     print(get_state_values(u))
90.     print(f_new)
91.  
92.  
93. if __name__ == '__main__':
94.     main()
95.