Update cosmotool 2nd part

external/cosmotool/python_sample/icgen/cosmogrowth.py

@@ -0,0 +1,202 @@
+import numexpr as ne
+import multiprocessing
+import pyfftw
+import weakref
+import numpy as np
+import cosmolopy as cpy
+import cosmotool as ct
+
+class CubeFT(object):
+  def __init__(self, L, N, max_cpu=-1):
+  
+      self.N = N
+      self.align = pyfftw.simd_alignment
+      self.L = L
+      self.max_cpu = multiprocessing.cpu_count() if max_cpu < 0 else max_cpu
+      self._dhat = pyfftw.n_byte_align_empty((self.N,self.N,self.N/2+1), self.align, dtype='complex64')
+      self._density = pyfftw.n_byte_align_empty((self.N,self.N,self.N), self.align, dtype='float32')
+      self._irfft = pyfftw.FFTW(self._dhat, self._density, axes=(0,1,2), direction='FFTW_BACKWARD', threads=self.max_cpu, normalize_idft=False)
+      self._rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False)
+
+  def rfft(self):
+      return ne.evaluate('c*a', local_dict={'c':self._rfft(normalise_idft=False),'a':(self.L/self.N)**3})
+      
+  def irfft(self):
+      return ne.evaluate('c*a', local_dict={'c':self._irfft(normalise_idft=False),'a':(1/self.L)**3})
+   
+  def get_dhat(self):
+      return self._dhat
+  def set_dhat(self, in_dhat):
+      self._dhat[:] = in_dhat 
+  dhat = property(get_dhat, set_dhat, None)
+      
+  def get_density(self):
+      return self._density      
+  def set_density(self, d):
+      self._density[:] = d 
+  density = property(get_density, set_density, None)
+      
+
+class CosmoGrowth(object):
+
+  def __init__(self, **cosmo):
+    self.cosmo = cosmo
+
+  def D(self, a):
+    return cpy.perturbation.fgrowth(1/a-1, self.cosmo['omega_M_0'], unnormed=True)
+
+  def compute_E(self, a):
+    om = self.cosmo['omega_M_0']
+    ol = self.cosmo['omega_lambda_0']
+    ok = self.cosmo['omega_k_0']
+
+    E = np.sqrt(om/a**3 + ol + ok/a**2)
+  
+    H2 = -3*om/a**4  - 2*ok/a**3
+
+    Eprime = 0.5*H2/E
+
+    return E,Eprime
+
+  def Ddot(self, a):
+    E,Eprime = self.compute_E(a)
+    D = self.D(a)
+    Ddot_D = Eprime/E + 2.5 * self.cosmo['omega_M_0']/(a**3*E**2*D)
+    Ddot_D *= a
+    return Ddot_D
+
+  def compute_velmul(self, a):
+    E,_ = self.compute_E(a)
+    velmul = self.Ddot(a)
+    velmul *= 100 * a * E
+    return velmul
+
+
+
+
+
+class LagrangianPerturbation(object):
+
+    def __init__(self,density,L, fourier=False, supersample=1, max_cpu=-1):
+    
+        self.L = L
+        self.N = density.shape[0]
+        
+        self.max_cpu = max_cpu
+        self.cube = CubeFT(self.L, self.N, max_cpu=max_cpu)
+        
+        if not fourier:
+          self.cube.density = density
+          self.dhat = self.cube.rfft().copy()
+        else:
+          self.dhat = density.copy()
+
+        if supersample > 1:
+          self.upgrade_sampling(supersample)
+        self.ik = np.fft.fftfreq(self.N, d=L/self.N)*2*np.pi
+        self._kx = self.ik[:,None,None]
+        self._ky = self.ik[None,:,None]
+        self._kz = self.ik[None,None,:(self.N/2+1)]
+        self.cache = {}#weakref.WeakValueDictionary()
+
+    @ct.timeit_quiet
+    def upgrade_sampling(self, supersample):
+        N2 = self.N * supersample
+        N = self.N
+        dhat_new = np.zeros((N2, N2, N2/2+1), dtype=np.complex128)
+
+        hN = N/2
+        dhat_new[:hN, :hN, :hN+1] = self.dhat[:hN, :hN, :]
+        dhat_new[:hN, (N2-hN):N2, :hN+1] = self.dhat[:hN, hN:, :]
+        dhat_new[(N2-hN):N2, (N2-hN):N2, :hN+1] = self.dhat[hN:, hN:, :]
+        dhat_new[(N2-hN):N2, :hN, :hN+1] = self.dhat[hN:, :hN, :]
+
+        self.dhat = dhat_new
+        self.N = N2
+        self.cube = CubeFT(self.L, self.N, max_cpu=self.max_cpu)
+        
+    @ct.timeit_quiet
+    def _gradient(self, phi, direction):
+        if direction == 'all':
+          dirs = [0,1,2]
+          copy = True
+        else:
+          dirs = [direction]
+          copy = False
+        ret=[]
+        for dir in dirs:
+          ne.evaluate('phi_hat * i * kv / (kx**2 + ky**2 + kz**2)', out=self.cube.dhat,
+              local_dict={'i':-1j, 'phi_hat':phi, 'kv':self._kdir(dir),
+                          'kx':self._kx, 'ky':self._ky, 'kz':self._kz},casting='unsafe')
+#        self.cube.dhat = self._kdir(direction)*1j*phi 
+          self.cube.dhat[0,0,0] = 0
+          x = self.cube.irfft()
+          ret.append(x.copy() if copy else x)
+        return ret[0] if len(ret)==1 else ret
+            
+    @ct.timeit_quiet
+    def lpt1(self, direction=0):
+        return self._gradient(self.dhat, direction)
+        
+    def new_shape(self,direction, q=3, half=False):
+        N0 = (self.N/2+1) if half else self.N
+        return ((1,)*direction) + (N0,) + ((1,)*(q-1-direction))
+
+    def _kdir(self, direction, q=3):
+        if direction != q-1:
+            return self.ik.reshape(self.new_shape(direction, q=q))
+        else:
+            return self.ik[:self.N/2+1].reshape(self.new_shape(direction, q=q, half=True))
+
+    def _get_k2(self, q=3):
+        if 'k2' in self.cache:
+            return self.cache['k2']
+            
+        k2 = self._kdir(0, q=q)**2
+        for d in xrange(1,q):
+            k2 = k2 + self._kdir(d, q=q)**2
+            
+        self.cache['k2'] = k2
+        return k2
+        
+    def _do_irfft(self, array, copy=True):
+        if copy:
+          self.cube.dhat = array
+        return self.cube.irfft()
+        
+    def _do_rfft(self, array, copy=True):
+        if copy:
+          self.cube.density = array
+        return self.cube.rfft()
+
+    @ct.timeit_quiet
+    def lpt2(self, direction=0):
+#        k2 = self._get_k2()
+#        k2[0,0,0] = 1
+
+        inv_k2 = ne.evaluate('1/(kx**2+ky**2+kz**2)', {'kx':self._kdir(0),'ky':self._kdir(1),'kz':self._kdir(2)})
+        inv_k2[0,0,0]=0
+        potgen0 = lambda i: ne.evaluate('kdir**2*d*ik2',out=self.cube.dhat,local_dict={'kdir':self._kdir(i),'d':self.dhat,'ik2':inv_k2}, casting='unsafe' )
+        potgen = lambda i,j: ne.evaluate('kdir0*kdir1*d*ik2',out=self.cube.dhat,local_dict={'kdir0':self._kdir(i),'kdir1':self._kdir(j),'d':self.dhat,'ik2':inv_k2}, casting='unsafe' )
+
+        if 'lpt2_potential' not in self.cache:
+            print("Rebuilding potential...")
+            div_phi2 = np.zeros((self.N,self.N,self.N), dtype=np.float64)
+            for j in xrange(3):
+                q = self._do_irfft( potgen0(j) ).copy()
+                for i in xrange(j+1, 3):
+                    with ct.time_block("LPT2 elemental (%d,%d)" %(i,j)):
+                      ne.evaluate('div + q * pot', out=div_phi2, 
+                            local_dict={'div':div_phi2, 'q':q,'pot':self._do_irfft( potgen0(i), copy=False ) } 
+                            )
+                      ne.evaluate('div - pot**2',out=div_phi2,
+                            local_dict={'div':div_phi2,'pot':self._do_irfft(potgen(i,j), copy=False) }
+                            )
+
+            phi2_hat = self._do_rfft(div_phi2)
+            #self.cache['lpt2_potential'] = phi2_hat
+            del div_phi2
+        else:
+            phi2_hat = self.cache['lpt2_potential']
+            
+        return self._gradient(phi2_hat, direction)