-
Notifications
You must be signed in to change notification settings - Fork 0
/
solver.py
153 lines (123 loc) · 5.76 KB
/
solver.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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
#################
# AA Solvers #
#################
import numpy as np
from itertools import product
from scipy.linalg import solve
from scipy.sparse.linalg import spsolve
def check_each_direction(n,angs,ifprint=True):
""" returns a list of the index of elements of n which do not have adequate
toy angle coverage. The criterion is that we must have at least one sample
in each Nyquist box when we project the toy angles along the vector n """
checks = np.array([])
P = np.array([])
if(ifprint):
print("\nChecking modes:\n====")
for k,i in enumerate(n):
N_matrix = np.linalg.norm(i)
X = np.dot(angs,i)
if(np.abs(np.max(X)-np.min(X))<2.*np.pi):
if(ifprint):
print "Need a longer integration window for mode ", i
checks=np.append(checks,i)
P = np.append(P,(2.*np.pi-np.abs(np.max(X)-np.min(X))))
elif(np.abs(np.max(X)-np.min(X))/len(X)>np.pi):
if(ifprint):
print "Need a finer sampling for mode ", i
checks=np.append(checks,i)
P = np.append(P,(2.*np.pi-np.abs(np.max(X)-np.min(X))))
if(ifprint):
print("====\n")
return checks,P
def solver(AA, N_max, symNx = 2, throw_out_modes=False):
""" Constructs the matrix A and the vector b from a timeseries of toy
action-angles AA to solve for the vector x = (J_0,J_1,J_2,S...) where
x contains all Fourier components of the generating function with |n|<N_max """
# Find all integer component n_vectors which lie within sphere of radius N_max
# Here we have assumed that the potential is symmetric x->-x, y->-y, z->-z
# This can be relaxed by changing symN to 1
# Additionally due to time reversal symmetry S_n = -S_-n so we only consider
# "half" of the n-vector-space
angs = unroll_angles(AA.T[3:].T,np.ones(3))
symNz = 2
NNx = range(-N_max, N_max+1, symNx)
NNy = range(-N_max, N_max+1, symNz)
NNz = range(-N_max, N_max+1, symNz)
n_vectors = np.array([[i,j,k] for (i,j,k) in product(NNx,NNy,NNz)
if(not(i==0 and j==0 and k==0) # exclude zero vector
and (k>0 # northern hemisphere
or (k==0 and j>0) # half of x-y plane
or (k==0 and j==0 and i>0)) # half of x axis
and np.sqrt(i*i+j*j+k*k)<=N_max)]) # inside sphere
check_each_direction(n_vectors,angs)
if(throw_out_modes):
n_vectors = np.delete(n_vectors,check_each_direction(n_vectors,angs),axis=0)
n = len(n_vectors)+3
b = np.zeros(shape=(n, ))
a = np.zeros(shape=(n,n))
a[:3,:3]=len(AA)*np.identity(3)
for i in AA:
a[:3,3:]+=2.*n_vectors.T[:3]*np.cos(np.dot(n_vectors,i[3:]))
a[3:,3:]+=4.*np.dot(n_vectors,n_vectors.T)*np.outer(np.cos(np.dot(n_vectors,i[3:])),np.cos(np.dot(n_vectors,i[3:])))
b[:3]+=i[:3]
b[3:]+=2.*np.dot(n_vectors,i[:3])*np.cos(np.dot(n_vectors,i[3:]))
a[3:,:3]=a[:3,3:].T
return np.array(solve(a,b)), n_vectors
from itertools import izip
def unroll_angles(A,sign):
""" Unrolls the angles, A, so they increase continuously """
n = np.array([0,0,0])
P = np.zeros(np.shape(A))
P[0]=A[0]
for i in xrange(1,len(A)):
n = n+((A[i]-A[i-1]+0.5*sign*np.pi)*sign<0)*np.ones(3)*2.*np.pi
P[i] = A[i]+sign*n
return P
import matplotlib.pyplot as plt
from scipy.stats import linregress as lr
def angle_solver(AA, timeseries, N_max, sign, symNx = 2, throw_out_modes=False):
""" Constructs the matrix A and the vector b from a timeseries of toy
action-angles AA to solve for the vector x = (theta_0,theta_1,theta_2,omega_1,
omega_2,omega_3, dSdx..., dSdy..., dSdz...) where x contains all derivatives
of the Fourier components of the generating function with |n| < N_max """
# First unroll angles
angs = unroll_angles(AA.T[3:].T,sign)
# Same considerations as above
symNz = 2
NNx = range(-N_max, N_max+1, symNx)
NNy = range(-N_max, N_max+1, symNz)
NNz = range(-N_max, N_max+1, symNz)
n_vectors = np.array([[i,j,k] for (i,j,k) in product(NNx,NNy,NNz)
if(not(i==0 and j==0 and k==0) # exclude zero vector
and (k>0 # northern hemisphere
or (k==0 and j>0) # half of x-y plane
or (k==0 and j==0 and i>0)) # half of x axis
and np.sqrt(i*i+j*j+k*k)<=N_max # inside sphere
)])
if(throw_out_modes):
n_vectors = np.delete(n_vectors,check_each_direction(n_vectors,angs),axis=0)
nv = len(n_vectors)
n = 3*nv+6
b = np.zeros(shape=(n, ))
a = np.zeros(shape=(n,n))
a[:3,:3]=len(AA)*np.identity(3)
a[:3,3:6]=np.sum(timeseries)*np.identity(3)
a[3:6,:3]=a[:3,3:6]
a[3:6,3:6]=np.sum(timeseries*timeseries)*np.identity(3)
for i,j in izip(angs,timeseries):
a[6:6+nv,0]+=-2.*np.sin(np.dot(n_vectors,i))
a[6:6+nv,3]+=-2.*j*np.sin(np.dot(n_vectors,i))
a[6:6+nv,6:6+nv]+=4.*np.outer(np.sin(np.dot(n_vectors,i)),np.sin(np.dot(n_vectors,i)))
b[:3]+=i
b[3:6]+=j*i
b[6:6+nv]+=-2.*i[0]*np.sin(np.dot(n_vectors,i))
b[6+nv:6+2*nv]+=-2.*i[1]*np.sin(np.dot(n_vectors,i))
b[6+2*nv:6+3*nv]+=-2.*i[2]*np.sin(np.dot(n_vectors,i))
a[6+nv:6+2*nv,1]=a[6:6+nv,0]
a[6+2*nv:6+3*nv,2]=a[6:6+nv,0]
a[6+nv:6+2*nv,4]=a[6:6+nv,3]
a[6+2*nv:6+3*nv,5]=a[6:6+nv,3]
a[6+nv:6+2*nv,6+nv:6+2*nv]=a[6:6+nv,6:6+nv]
a[6+2*nv:6+3*nv,6+2*nv:6+3*nv]=a[6:6+nv,6:6+nv]
a[:6,:]=a[:,:6].T
return np.array(solve(a,b))