-
Notifications
You must be signed in to change notification settings - Fork 10
/
balance.py
49 lines (40 loc) · 1.22 KB
/
balance.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
import numpy as np
from openopt import NLP
np.set_printoptions(suppress=True)
def balance(sam, debug=False):
try:
table = sam.array()
except AttributeError:
table = np.array(sam)
assert table.shape[0] == table.shape[1]
size = table.shape[0]
x0 = np.array([v for v in table.flatten() if v !=0])
def transform(ox):
ret = np.zeros_like(table)
i = 0
for r in range(size):
for c in range(size):
if table[r, c] != 0:
ret[r, c] = ox[i]
i += 1
return ret
def objective(ox):
ox = np.square((ox - x0) / x0)
return np.sum(ox)
def constraints(ox):
ox = transform(ox)
ret = np.sum(ox, 0) - np.sum(ox, 1)
return ret
print constraints(x0)
if debug:
print("--- balance ---")
p = NLP(objective, x0, h=constraints, iprint = 50 * int(debug), maxIter = 100000, maxFunEvals = 1e7, name = 'NLP_1')
r = p.solve('ralg', plot=0)
if debug:
print 'constraints'
print constraints(r.xf)
assert r.isFeasible
try:
return sam.replace(transform(r.xf))
except UnboundLocalError:
return transform(r.xf)