import numpy as np
from experiment import graph_cochains, select_matroid, residual, train_models, wrap
np.random.seed(573)
N=900; t=np.random.rand(N,2)
x=np.c_[np.cos(2*np.pi*t[:,0]),np.sin(2*np.pi*t[:,0]),np.cos(2*np.pi*t[:,1]),np.sin(2*np.pi*t[:,1])]
theta=np.c_[t[:,0],t[:,1],(t[:,0]+t[:,1])%1,(t[:,0]-t[:,1])%1]
G,A,e=graph_cochains(theta)
Q=np.array([[1,0],[0,1],[.5,.5],[.5,-.5]],float)
sel,E,M=select_matroid(G,Q)
print({'selected':sel,'energies':E.tolist(),'weight_sum':1.0,'gram_symmetry':float(np.max(abs(G-G.T)))})
print({'residual_full':residual(G,Q,[0,1])[0],'residual_omit_one':residual(G,Q,[0])[0]})
print({'train':train_models(x,sel,theta)})