Coverage for python/lsst/sims/featureScheduler/thomson/thomson.py : 8%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import numpy as np
2from scipy.optimize import minimize
3from lsst.sims.utils import _angularSeparation
6__all__ = ['thetaphi2xyz', 'even_points', 'elec_potential', 'ang_potential', 'fib_sphere_grid',
7 'iterate_potential_random', 'iterate_potential_smart', 'even_points_xyz', 'elec_potential_xyz',
8 'xyz2thetaphi']
11def thetaphi2xyz(theta, phi):
12 x = np.sin(phi)*np.cos(theta)
13 y = np.sin(phi)*np.sin(theta)
14 z = np.cos(phi)
15 return x, y, z
18def xyz2thetaphi(x, y, z):
19 phi = np.arccos(z)
20 theta = np.arctan2(y, x)
21 return theta, phi
24def elec_potential(x0):
25 """
26 Compute the potential energy for electrons on a sphere
28 Parameters
29 ----------
30 x0 : array
31 First half of x0 or theta values, secnd half phi
33 Returns
34 -------
35 Potential energy
36 """
38 theta = x0[0:int(x0.size/2)]
39 phi = x0[int(x0.size/2):]
41 x, y, z = thetaphi2xyz(theta, phi)
42 # Distance squared
43 dsq = 0.
45 indices = np.triu_indices(x.size, k=1)
47 for coord in [x, y, z]:
48 coord_i = np.tile(coord, (coord.size, 1))
49 coord_j = coord_i.T
50 d = (coord_i[indices]-coord_j[indices])**2
51 dsq += d
53 U = np.sum(1./np.sqrt(dsq))
54 return U
57def potential_single(coord0, x, y, z):
58 """
59 Find the potential contribution from a single point.
60 """
62 x0 = coord0[0]
63 y0 = coord0[1]
64 z0 = coord0[2]
65 # Enforce point has to be on a sphere
66 rsq = x0**2+y0**2 + z0**2
67 r = np.sqrt(rsq)
68 x0 = x0/r
69 y0 = y0/r
70 z0 = z0/r
72 dsq = (x-x0)**2+(y-y0)**2+(z-z0)**2
73 U = np.sum(1./np.sqrt(dsq))
74 return U
77def xyz2U(x, y, z):
78 """
79 compute the potential
80 """
81 dsq = 0.
83 indices = np.triu_indices(x.size, k=1)
85 for coord in [x, y, z]:
86 coord_i = np.tile(coord, (coord.size, 1))
87 coord_j = coord_i.T
88 dsq += (coord_i[indices]-coord_j[indices])**2
90 d = np.sqrt(dsq)
91 U = np.sum(1./d)
92 return U
95def iterate_potential_smart(x0, stepfrac=0.1):
96 """
97 Calculate the change in potential by shifting points in theta and phi directions
98 # wow, that sure didn't work at all.
99 """
101 theta = x0[0:x0.size/2]
102 phi = x0[x0.size/2:]
103 x, y, z = thetaphi2xyz(theta, phi)
104 U_input = xyz2U(x, y, z)
106 # Now to loop over each point, and find where it's potenital minimum would be, and move it
107 # half-way there.
108 xyz_new = np.zeros((x.size, 3), dtype=float)
109 mask = np.ones(x.size, dtype=bool)
110 for i in np.arange(x.size):
111 mask[i] = 0
112 fit = minimize(potential_single, [x[i], y[i], z[i]], args=(x[mask], y[mask], z[mask]))
113 mask[i] = 1
114 xyz_new[i] = fit.x/np.sqrt(np.sum(fit.x**2))
116 xyz_input = np.array((x, y, z)).T
117 diff = xyz_input - xyz_new
119 # Move half way in x-y-z space
120 xyz_out = xyz_input + stepfrac*diff
121 # Project back onto sphere
122 xyz_out = xyz_out.T/np.sqrt(np.sum(xyz_out**2, axis=1))
123 U_new = xyz2U(xyz_out[0, :], xyz_out[1, :], xyz_out[2, :])
124 theta, phi = xyz2thetaphi(xyz_out[0, :], xyz_out[1, :], xyz_out[2, :])
125 return np.concatenate((theta, phi)), U_new
128def iterate_potential_random(x0, stepsize=.05):
129 """
130 Given a bunch of theta,phi values, shift things around to minimize potential
131 """
133 theta = x0[0:int(x0.size/2)]
134 phi = x0[int(x0.size/2):]
136 x, y, z = thetaphi2xyz(theta, phi)
137 # Distance squared
138 dsq = 0.
140 indices = np.triu_indices(x.size, k=1)
142 for coord in [x, y, z]:
143 coord_i = np.tile(coord, (coord.size, 1))
144 coord_j = coord_i.T
145 d = (coord_i[indices]-coord_j[indices])**2
146 dsq += d
148 d = np.sqrt(dsq)
150 U_input = 1./d
152 # offset everything by a random ammount
153 x_new = x + np.random.random(theta.size) * stepsize
154 y_new = y + np.random.random(theta.size) * stepsize
155 z_new = z + np.random.random(theta.size) * stepsize
157 r = (x_new**2 + y_new**2 + z_new**2)**0.5
158 # put back on the sphere
159 x_new = x_new/r
160 y_new = y_new/r
161 z_new = z_new/r
163 dsq_new = 0
164 for coord, coord_new in zip([x, y, z], [x_new, y_new, z_new]):
165 coord_i_new = np.tile(coord_new, (coord_new.size, 1))
166 coord_j = coord_i_new.T
167 d_new = (coord_i_new[indices]-coord_j[indices])**2
168 dsq_new += d_new
169 U_new = 1./np.sqrt(dsq_new)
171 U_diff = np.sum(U_new)-np.sum(U_input)
172 if U_diff > 0:
173 return x0, 0.
174 else:
175 theta, phi = xyz2thetaphi(x_new, y_new, z_new)
176 return np.concatenate((theta, phi)), U_diff
179def ang_potential(x0):
180 """
181 If distance is computed along sphere rather than through 3-space.
182 """
183 theta = x0[0:int(x0.size/2)]
184 phi = np.pi/2-x0[int(x0.size/2):]
186 indices = np.triu_indices(theta.size, k=1)
188 theta_i = np.tile(theta, (theta.size, 1))
189 theta_j = theta_i.T
190 phi_i = np.tile(phi, (phi.size, 1))
191 phi_j = phi_i.T
192 d = _angularSeparation(theta_i[indices], phi_i[indices], theta_j[indices], phi_j[indices])
193 U = np.sum(1./d)
194 return U
197def fib_sphere_grid(npoints):
198 """
199 Use a Fibonacci spiral to distribute points uniformly on a sphere.
201 based on https://people.sc.fsu.edu/~jburkardt/py_src/sphere_fibonacci_grid/sphere_fibonacci_grid_points.py
203 Returns theta and phi in radians
204 """
206 phi = (1.0 + np.sqrt(5.0)) / 2.0
208 i = np.arange(npoints, dtype=float)
209 i2 = 2*i - (npoints-1)
210 theta = (2.0*np.pi * i2/phi) % (2.*np.pi)
211 sphi = i2/npoints
212 phi = np.arccos(sphi)
213 return theta, phi
216def even_points(npts, use_fib_init=True, method='CG', potential_func=elec_potential, maxiter=None):
217 """
218 Distribute npts over a sphere and minimize their potential, making them
219 "evenly" distributed
221 Starting with the Fibonacci spiral speeds things up by ~factor of 2.
222 """
224 if use_fib_init:
225 # Start with fibonacci spiral guess
226 theta, phi = fib_sphere_grid(npts)
227 else:
228 # Random on a sphere
229 theta = np.random.rand(npts)*np.pi*2.
230 phi = np.arccos(2.*np.random.rand(npts)-1.)
232 x = np.concatenate((theta, phi))
233 # XXX--need to check if this is the best minimizer
234 min_fit = minimize(potential_func, x, method='CG', options={'maxiter': maxiter})
236 x = min_fit.x
237 theta = x[0:int(x.size/2)]
238 phi = x[int(x.size/2):]
239 # Looks like I get the same energy values as https://en.wikipedia.org/wiki/Thomson_problem
240 return theta, phi
243def elec_potential_xyz(x0):
244 x0 = x0.reshape(3, int(x0.size/3))
245 x = x0[0, :]
246 y = x0[1, :]
247 z = x0[2, :]
248 dsq = 0.
250 r = np.sqrt(x**2 + y**2 + z**2)
251 x = x/r
252 y = y/r
253 z = z/r
254 indices = np.triu_indices(x.size, k=1)
256 for coord in [x, y, z]:
257 coord_i = np.tile(coord, (coord.size, 1))
258 coord_j = coord_i.T
259 d = (coord_i[indices]-coord_j[indices])**2
260 dsq += d
262 U = np.sum(1./np.sqrt(dsq))
263 return U
266def elec_p_xyx_loop(x0):
267 """do this with a brutal loop that can be numba ified
268 """
269 x0 = x0.reshape(3, int(x0.size/3))
270 x = x0[0, :]
271 y = x0[1, :]
272 z = x0[2, :]
273 U = 0.
275 r = np.sqrt(x**2 + y**2 + z**2)
276 x = x/r
277 y = y/r
278 z = z/r
280 npts = x.size
281 for i in range(npts-1):
282 for j in range(i+1, npts):
283 dsq = (x[i]-x[j])**2 + (y[i]-y[j])**2 + (z[i]-z[j])**2
284 U += 1./np.sqrt(dsq)
285 return U
288def x02sphere(x0):
289 x0 = x0.reshape(3, int(x0.size/3))
290 x = x0[0, :]
291 y = x0[1, :]
292 z = x0[2, :]
294 r = np.sqrt(x**2 + y**2 + z**2)
295 x = x/r
296 y = y/r
297 z = z/r
299 return np.concatenate((x, y, z))
302def even_points_xyz(npts, use_fib_init=True, method='CG', potential_func=elec_p_xyx_loop, maxiter=None,
303 callback=None, verbose=True):
304 """
305 Distribute npts over a sphere and minimize their potential, making them
306 "evenly" distributed
308 Starting with the Fibonacci spiral speeds things up by ~factor of 2.
309 """
311 if use_fib_init:
312 # Start with fibonacci spiral guess
313 theta, phi = fib_sphere_grid(npts)
314 else:
315 # Random on a sphere
316 theta = np.random.rand(npts)*np.pi*2.
317 phi = np.arccos(2.*np.random.rand(npts)-1.)
319 x = np.concatenate(thetaphi2xyz(theta, phi))
321 if verbose:
322 print('initial potential=', elec_potential_xyz(x))
323 # XXX--need to check if this is the best minimizer
324 min_fit = minimize(potential_func, x, method='CG', options={'maxiter': maxiter}, callback=callback)
326 if verbose:
327 print('final potential=', elec_potential_xyz(min_fit.x))
328 print('niteration=', min_fit.nit)
330 x = x02sphere(min_fit.x)
332 # Looks like I get the same energy values as https://en.wikipedia.org/wiki/Thomson_problem
333 return x