forked from neyrinck/FoldYourOwnUniverse
-
Notifications
You must be signed in to change notification settings - Fork 0
/
foldyourown.py
109 lines (83 loc) · 2.8 KB
/
foldyourown.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
import sys,os
sys.path.append(os.path.realpath('..'))
import numpy as N, pylab as M
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
# These variables are global, so they're defined up here first
scale = 1.0
slider_scale = 1.0
psi = 1.0
def init():
""" Set up the initial conditions """
# Sets up initial dimensions
M.figure(figsize=(8,8))
density=N.loadtxt(os.path.join(sys.path[0], './densmesh.txt'))
# Real Fourier transform of "density"
density_k = N.fft.rfftn(density*1e3) # 1e3 to enhance contrast
global psi
psi = zeldovich(density_k)
# Zel'dovich displacement field
global scale
global slider_scale
scale = 1.0
slider_scale = 50.0
axcolor = 'lightgoldenrodyellow'
axScale = plt.axes([0.15, 0.1, 0.7, 0.03], axisbg=axcolor)
slider_scale = Slider(axScale, 'Scale', 0.0, 50.0, 1.0)
slider_scale.on_changed(update)
# Attempting to removed axes from the graph
plt.axes([0.15, 0.15, 0.7, 0.7])
plt.xticks([])
plt.yticks([])
return density_k, psi
def getkgrid(ng=64):
""" ng = number of particles in each dim """
thirdim=ng/2+1
sk = (ng,thirdim)
kx = N.fromfunction(lambda x,y:x, sk).astype(N.float32)
kx[N.where(kx > ng/2)] -= ng
ky = N.fromfunction(lambda x,y:y, sk).astype(N.float32)
ky[N.where(ky > ng/2)] -= ng
k2 = kx**2+ky**2
k2[0,0] = 1.
return kx/k2, ky/k2
def zeldovich(dk,boxsize=None):
""" implements the Zel'dovich approximation to move particles """
sk = dk.shape #shape of the Fourier-space density field
ng = sk[0] #number of particles in each dimension
kx, ky = getkgrid(ng)
# psi = displacement of each particle from initial conditions.
psi = N.empty((ng,ng,2),dtype=N.float32)
psi[:,:,0] = N.fft.irfftn(-1j*dk*kx)
psi[:,:,1] = N.fft.irfftn(-1j*dk*ky)
return psi
def psi2pos(psi,boxsize=500.):
"""
displacement field to positions
(adding displacement to a regular lattice)
"""
ng = psi.shape[0] # number of particles in each dimension
x=boxsize/ng*N.fromfunction(lambda x,y:x, (ng,ng))
# x positions on a regular lattice
pos = 1.*psi
pos[:,:,0] += x
pos[:,:,1] += N.transpose(x)
return pos
def plotvertices(psi):
boxsize = 500.
pos = psi2pos(psi*scale,boxsize)
plt.axes([0.15, 0.15, 0.7, 0.7])
plt.xticks([])
plt.yticks([])
plt.xlim((boxsize * -0.2), (boxsize * 1.2))
plt.ylim((boxsize * -0.2), (boxsize * 1.2))
#plot the vertices
M.scatter(pos[:,:,0].flat,pos[:,:,1].flat,s=1,lw=0)
# When the slider is changed, redraw the screen with the updated scale
def update(val):
global psi
global scale
global slider_scale
scale = slider_scale.val
M.cla()
plotvertices(psi)