Moved code. Reorganized more OO like

python_sample/icgen/cosmogrowth.py

@@ -0,0 +1,102 @@
+import weakref
+import numpy as np
+import cosmolopy as cpy
+
+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):
+    
+        self.L = L
+        self.N = density.shape[0]
+        self.dhat = np.fft.rfftn(density)*(L/self.N)**3 if not fourier else density
+        self.ik = np.fft.fftfreq(self.N, d=L/self.N)*2*np.pi
+        self.cache = weakref.WeakValueDictionary()
+        
+    def _gradient(self, phi, direction):
+        return np.fft.irfftn(self._kdir(direction)*1j*phi)*(self.N/self.L)**3
+            
+    def lpt1(self, direction=0):
+        k2 = self._get_k2()
+        k2[0,0,0] = 1
+    
+        return self._gradient(-self.dhat/k2, 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 lpt2(self, direction=0):
+        k2 = self._get_k2()
+        k2[0,0,0] = 1
+
+        if 'lpt2_potential' not in self.cache:
+            div_phi2 = np.zeros((N,N,N), dtype=np.float64)
+            for j in xrange(3):
+                q = np.fft.irfftn( build_dir(ik, j)**2*self.dhat / k2 )
+                for i in xrange(j+1, 3):
+                    div_phi2 += q * np.fft.irfftn( build_dir(ik, i)**2*self.dhat / k2 )
+                    div_phi2 -= (np.fft.irfftn( build_dir(ik, j)*build_dir(ik, i)*self.dhat / k2 ))**2
+
+            div_phi2 *= (self.N/self.L)**3
+            phi2_hat = np.fft.rfftn(div_phi2) * ((L/N)**3) / k2
+            self.cache['lpt2_potential'] = phi2_hat
+            del div_phi2
+        else:
+            phi2_hat = self.cache['lpt2_potential']
+            
+        return self._gradient(phi2_hat, direction)