1import numpy as np
 2from router import exact_stats, sample_subset_dp, natural_direction
 3
 4theta=np.array([2.0,1.0,0.2,-0.4,-1.0,-2.0])
 5m=2
 6mu,sigma=exact_stats(theta,m)
 7rng=np.random.default_rng(11)
 8counts=np.zeros(len(theta))
 9N=200000
10for _ in range(N):
11    counts[sample_subset_dp(theta,m,rng)] += 1
12emp=counts/N
13err=np.max(np.abs(emp-mu))
14g=np.array([1.2,-.7,.4,.1,-.5,-.5])
15d,info=natural_direction(theta,m,g,rho=.17)
16v=np.diag(info['sigma'])
17cap=max(abs(d[i]-d[j])/np.sqrt(1/v[i]+1/v[j]) for i in range(len(theta)) for j in range(i+1,len(theta)))
18print({'exact_mu':mu.tolist(),'empirical_mu':emp.tolist(),'max_sampling_error':float(err), 'sum_direction':float(d.sum()), 'cap':float(cap),'rho':.17,'raw_cap':info['raw_cap']})
19assert err < .006
20assert abs(d.sum()) < 1e-10
21assert cap <= .17 + 1e-10