generate ksample datasets

bfps/fluid_base.py

@@ -50,15 +50,15 @@ class fluid_particle_base(bfps.code):
             self.dtype = dtype
         elif dtype in ['single', 'double']:
             if dtype == 'single':
-                self.dtype = np.float32
+                self.dtype = np.dtype(np.float32)
             elif dtype == 'double':
-                self.dtype = np.float64
+                self.dtype = np.dtype(np.float64)
         self.rtype = self.dtype
         if self.rtype == np.float32:
-            self.ctype = np.complex64
+            self.ctype = np.dtype(np.complex64)
             self.C_dtype = 'float'
         elif self.rtype == np.float64:
-            self.ctype = np.complex128
+            self.ctype = np.dtype(np.complex128)
             self.C_dtype = 'double'
         self.parameters['dkx'] = 1.0
         self.parameters['dky'] = 1.0
@@ -275,19 +275,19 @@ class fluid_particle_base(bfps.code):
             write_to_file = False):
         np.random.seed(rseed)
         Kdata00 = bfps.tools.generate_data_3D(
-                self.parameters['nz']/2,
-                self.parameters['ny']/2,
-                self.parameters['nx']/2,
+                self.parameters['nz']//2,
+                self.parameters['ny']//2,
+                self.parameters['nx']//2,
                 p = spectra_slope).astype(self.ctype)
         Kdata01 = bfps.tools.generate_data_3D(
-                self.parameters['nz']/2,
-                self.parameters['ny']/2,
-                self.parameters['nx']/2,
+                self.parameters['nz']//2,
+                self.parameters['ny']//2,
+                self.parameters['nx']//2,
                 p = spectra_slope).astype(self.ctype)
         Kdata02 = bfps.tools.generate_data_3D(
-                self.parameters['nz']/2,
-                self.parameters['ny']/2,
-                self.parameters['nx']/2,
+                self.parameters['nz']//2,
+                self.parameters['ny']//2,
+                self.parameters['nx']//2,
                 p = spectra_slope).astype(self.ctype)
         Kdata0 = np.zeros(
                 Kdata00.shape + (3,),
@@ -376,6 +376,8 @@ class fluid_particle_base(bfps.code):
         kspace['kz'] = np.arange(-self.parameters['nz']//2 + 1,
                                   self.parameters['nz']//2 + 1).astype(np.float64)*self.parameters['dkz']
         kspace['kz'] = np.roll(kspace['kz'], self.parameters['nz']//2+1)
+        kspace['ksample_indices'] = bfps.tools.get_kindices(n = self.parameters['nx'])
+        print kspace['ksample_indices'].shape[0]
         return kspace
     def write_par(self, iter0 = 0):
         assert (self.parameters['niter_todo'] % self.parameters['niter_stat'] == 0)
@@ -390,6 +392,14 @@ class fluid_particle_base(bfps.code):
             for k in kspace.keys():
                 ofile['kspace/' + k] = kspace[k]
             nshells = kspace['nshell'].shape[0]
+            time_chunk = 2**20//(self.ctype.itemsize*3*kspace['ksample_indices'].shape[0])
+            time_chunk = max(time_chunk, 1)
+            ofile.create_dataset('statistics/ksamples/velocity',
+                                 (1, kspace['ksample_indices'].shape[0], 3),
+                                 chunks = (time_chunk, kspace['ksample_indices'].shape[0], 3),
+                                 maxshape = (None, kspace['ksample_indices'].shape[0], 3),
+                                 dtype = self.ctype,
+                                 compression = 'gzip')
             for k in ['velocity', 'vorticity']:
                 time_chunk = 2**20//(8*3*self.parameters['nx'])
                 time_chunk = max(time_chunk, 1)

bfps/tools.py

@@ -24,6 +24,7 @@
 
 
 
+import math
 import numpy as np
 
 def generate_data_3D(
@@ -49,7 +50,7 @@ def generate_data_3D(
         a[0, n0-kz, 0] = np.conj(a[0, kz, 0])
     for ky in range(1, n1//2):
         for kz in range(1, n0):
-            a[n-ky, n-kz, 0] = np.conj(a[ky, kz, 0])
+            a[n1-ky, n0-kz, 0] = np.conj(a[ky, kz, 0])
     ii = np.where(k > min(n0, n1, n2)/3.)
     a[ii] = 0
     return a
@@ -73,6 +74,45 @@ def padd_with_zeros(
     b[n0-m0/2:    , n1-m1/2:    , :m2/2+1] = a[m0-m0/2:    , m1-m1/2:    , :m2/2+1]
     return b
 
+def get_kindices(
+        n = 64):
+    nx = n
+    nz = n
+    kx = np.arange(0, nx//2+1, 1).astype(np.float)
+    kvals = []
+    radii = set([])
+    index = []
+    for iz in range(kx.shape[0]):
+        for ix in range(0, kx.shape[0]):
+            kval = (kx[iz]**2+kx[ix]**2)**.5
+            tmp = math.modf(kval)
+            if (tmp[0] == 0 and tmp[1] <= nx//2):
+                kvals.append([kx[iz], kx[ix]])
+                radii.add(math.floor(kval))
+                index.append([ix, iz])
+
+    kvals = np.array(kvals)
+    index = np.array(index)
+    new_kvals = []
+    ordered_kvals = []
+    radius_vals = []
+    ii = []
+    for r in radii:
+        ncircle = np.count_nonzero((kvals[:, 0]**2 + kvals[:, 1]**2)**.5 == r)
+        indices = np.where((kvals[:, 0]**2 + kvals[:, 1]**2)**.5 == r)[0]
+        if ncircle > 2:
+            ordered_kvals.append(kvals[indices])
+            new_kvals += list(kvals[indices])
+            radius_vals.append(r)
+            ii += list(index[indices])
+    ii = np.array(ii)
+    good_indices = np.where(ii[:, 1] > 0)[0]
+    i1 = (kx.shape[0] - 1)*2 - ii[good_indices, 1]
+    i0 = ii[good_indices, 0]
+    i1 = np.concatenate((ii[:, 1], i1)),
+    i0 = np.concatenate((ii[:, 0], i0)),
+    return np.vstack((i0, i1)).T.copy()
+
 try:
     import sympy as sp
     rational0 = sp.Rational(0)