Adding possibility to eliminate Nyquist frequency

conf/conda/meta.yaml

@@ -6,7 +6,7 @@ source:
   git_url: ../../
 
 build:
-  number: 24
+  number: 25
 
 requirements:
   build:

mpiFFT4py/line.py

@@ -109,16 +109,23 @@ def get_local_mesh(self):
         X[1] *= self.L[1]/self.N[1]
         return X
 
-    def get_local_wavenumbermesh(self, scaled=True):
+    def get_local_wavenumbermesh(self, scaled=True, broadcast=False,
+                                 eliminate_highest_freq=False):
         kx = fftfreq(self.N[0], 1./self.N[0])
-        ky = fftfreq(self.N[1], 1./self.N[1])[:self.Nf]
-        ky[-1] *= -1
+        ky = rfftfreq(self.N[1], 1./self.N[1])
+        if eliminate_highest_freq:
+            for i, k in enumerate((kx, ky)):
+                if self.N[i] % 2 == 0:
+                    k[self.N[i]//2] = 0
+
         Ks = np.meshgrid(kx, ky[self.rank*self.Np[1]//2:(self.rank*self.Np[1]//2+self.Npf)], indexing='ij', sparse=True)
         if scaled is True:
             Lp = 2*np.pi/self.L
             Ks[0] *= Lp[0]
             Ks[1] *= Lp[1]
-        K = [np.broadcast_to(k, self.complex_shape()) for k in Ks]
+        K = Ks
+        if broadcast is True:
+            K = [np.broadcast_to(k, self.complex_shape()) for k in Ks]
         return K
 
     def get_dealias_filter(self):

mpiFFT4py/pencil.py

@@ -309,20 +309,34 @@ def get_local_mesh(self):
         X = [np.broadcast_to(x, self.real_shape()) for x in X]
         return X
 
-    def get_local_wavenumbermesh(self, scaled=False):
+    def get_local_wavenumbermesh(self, scaled=False, broadcast=False,
+                                 eliminate_highest_freq=False):
         """Returns (scaled) local decomposed wavenumbermesh
 
         If scaled is True, then the wavenumbermesh is scaled with physical mesh
         size. This takes care of mapping the physical domain to a computational
         cube of size (2pi)**3
+
+
         """
-        kx, ky, kz = self.complex_local_wavenumbers()
+        s = self.complex_local_slice()
+        kx = fftfreq(self.N[0], 1./self.N[0]).astype(int)
+        ky = fftfreq(self.N[1], 1./self.N[1]).astype(int)
+        kz = rfftfreq(self.N[2], 1./self.N[2]).astype(int)
+        if eliminate_highest_freq:
+            for i, k in enumerate((kx, ky, kz)):
+                if self.N[i] % 2 == 0:
+                    k[self.N[i]//2] = 0
+        kx = kx[s[0]]
+        kz = kz[s[2]]
         Ks = np.meshgrid(kx, ky, kz, indexing='ij', sparse=True)
         if scaled is True:
             Lp = 2*np.pi/self.L
             for i in range(3):
                 Ks[i] = (Ks[i]*Lp[i]).astype(self.float)
-        K = [np.broadcast_to(k, self.complex_shape()) for k in Ks]
+        K = Ks
+        if broadcast is True:
+            K = [np.broadcast_to(k, self.complex_shape()) for k in Ks]
         return K
 
     def get_dealias_filter(self):

mpiFFT4py/slab.py

@@ -159,21 +159,31 @@ def get_local_mesh(self):
         X = [np.broadcast_to(x, self.real_shape()) for x in X]
         return X
 
-    def get_local_wavenumbermesh(self, scaled=False):
+    def get_local_wavenumbermesh(self, scaled=False, broadcast=False, eliminate_highest_freq=False):
         """Returns (scaled) local decomposed wavenumbermesh
 
         If scaled is True, then the wavenumbermesh is scaled with physical mesh
         size. This takes care of mapping the physical domain to a computational
-        cube of size (2pi)**3
+        cube of size (2pi)**3.
+
+        If eliminate_highest_freq is True, then the Nyquist frequency is set to zero.
         """
         kx, ky, kz = self.complex_local_wavenumbers()
+        if eliminate_highest_freq:
+            ky = fftfreq(self.N[1], 1./self.N[1])
+            for i, k in enumerate((kx, ky, kz)):
+                if self.N[i] % 2 == 0:
+                    k[self.N[i]//2] = 0
+            ky = ky[self.complex_local_slice()[1]]
+
         Ks = np.meshgrid(kx, ky, kz, indexing='ij', sparse=True)
         if scaled:
             Lp = 2*np.pi/self.L
             for i in range(3):
                 Ks[i] *= Lp[i]
-        K = [np.broadcast_to(k, self.complex_shape()) for k in Ks]
-        #K = np.array(np.meshgrid(kx, ky, kz, indexing='ij')).astype(self.float)
+        K = Ks
+        if broadcast is True:
+            K = [np.broadcast_to(k, self.complex_shape()) for k in Ks]
         return K
 
     def get_dealias_filter(self):