mpi4py / mpi4py-fft

mpi4py-fft is a Python package for computing Fast Fourier Transforms (FFTs).
Other
50 stars 12 forks source link

switching to real to real transforms #11

Closed francispoulin closed 4 years ago

francispoulin commented 4 years ago

Hello Mikael,

I have figured out how to go between fftn/ifftn and rfftn/irfftn and I think this is working nicely.

I am now trying to figure out how to use cosine and sine transforms. My first attempt seemed to crash and burn. I defined the following but it seems to keep things periodic at the boundaries.

from mpi4py_fft.fftw import dstn, idstn, dctn, idctn

import functools
dst = functools.partial(dstn, type=3)
idst = functools.partial(idstn, type=3)

r2r = PFFT(MPI.COMM_WORLD, N, axes=(0, 1), dtype=np.float,
           grid=(-1,), transforms={(0,1): (dstn, idstn)} )

I also used the same wavenumbers that I used for the rfftn. Is there anything else that I should be doing?

mikaem commented 4 years ago

I don't quite understand the problem. dst/idst simply compute the transforms. I don't think it keeps things periodic?

Wavenumbers for r2r are simply arange(N), if N is an integer. Note that if you want to use type 3 above, then you need to use (dst, idst) in the creation of PFFT.

francispoulin commented 4 years ago

Sorry I was unclear. I don't want dst/idst to keep things periodic, but my simulation seems to be doing that, which makes me think that I'm not doing the correct transform.

To try and isolate the problem, using dct instead of dst for no good reason, and below is an example code that computes the dct in each direction and computes the second derivative in x, which is still a cosine function, and I'm getting a strange (incorrect) result. Am I doing something wrong?

`# dct2_example.py

import numpy as np
import matplotlib.pyplot as plt

from mpi4py import MPI
from mpi4py_fft import PFFT, newDistArray
from mpi4py_fft.fftw import rfftn, irfftn, dctn, idctn
import functools
dct = functools.partial(dctn, type=3)
idct = functools.partial(idctn, type=3)
transforms = {(0,1):  (dct, idct)}

def get_local_mesh(FFT, L):
    """Returns local mesh."""
    X = np.ogrid[FFT.local_slice(False)]
    N = FFT.global_shape()
    for i in range(len(N)):
        X[i] = (X[i]*L[i]/N[i])
    X = [np.broadcast_to(x, FFT.shape(False)) for x in X]
    return X

def get_local_wavenumbermesh(FFT, L):
    """Returns local wavenumber mesh."""
    s = FFT.local_slice()
    N = FFT.global_shape()
    # Set wavenumbers in grid
    k = [np.fft.fftfreq(n, 1./n).astype(int) for n in N[:-1]]
    k.append(np.fft.rfftfreq(N[-1], 1./N[-1]).astype(int))
    K = [ki[si] for ki, si in zip(k, s)]
    Ks = np.meshgrid(*K, indexing='ij', sparse=True)
    Lp = 2*np.pi/L
    for i in range(2):
        Ks[i] = (Ks[i]*Lp[i]).astype(float)
    return [np.broadcast_to(k, FFT.shape(True)) for k in Ks]

M = 6
N = [2**M, 2**M]
L = np.array([np.pi, np.pi], dtype=float) 

r2r = PFFT(MPI.COMM_WORLD, N, transforms=transforms)

X = get_local_mesh(r2r, L)
K = get_local_wavenumbermesh(r2r, L)
K = np.array(K).astype(float)
K2 = np.sum(K*K, 0, dtype=float)

u = newDistArray(r2r, False)
u = np.cos(X[0])*np.cos(X[1])
u_hat = -K[0]*K[0]*r2r.forward(u)
#u_hat = r2r.forward(u)
uj = np.zeros_like(u)
uj = r2r.backward(u_hat, uj)

plt.pcolormesh(u)
plt.colorbar()
plt.title('u')
plt.show()
plt.pcolormesh(uj)
plt.colorbar()
plt.title('uj')
plt.show()

`

mikaem commented 4 years ago

You do not have any Fourier transforms here. You use dct/idct for both axes. So you should not use fftfreq or rfftfreq. If you want one direction with Fourier and the other with cosine, then use transforms={0: (rfftn, irfftn), 1: (dct, idct)}. Or use Shenfun like for the Poisson equation here, which is using a Chebyshev basis along the first direction and Fourier along the second.

francispoulin commented 4 years ago

Thanks. Just so I understand, if I want to use dct in both directions I shouldn't use fftfreq or rfftfreq? Does that mean it is possible to do a double dct if I set up my wavenumbers correctly?

In the past, what I have done to compute the double dct, or dct2, is to extend the function evenly in both directions (4 times as much storage unfortunately) then use the fft2 to differentiate on the extended domain, then invert and project onto the physical domain.

This is something I'm trying to avoid in part because of the storage but also because I don't know how I would parallelize this projection, and probably don't want to. But it has the advantage of being able to go between sine and cosine transforms. Any advice?

mikaem commented 4 years ago

Thanks. Just so I understand, if I want to use dct in both directions I shouldn't use fftfreq or rfftfreq? Does that mean it is possible to do a double dct if I set up my wavenumbers correctly?

Yes, you can do a double dct.

In the past, what I have done to compute the double dct, or dct2, is to extend the function evenly in both directions (4 times as much storage unfortunately) then use the fft2 to differentiate on the extended domain, then invert and project onto the physical domain.

This is something I'm trying to avoid in part because of the storage but also because I don't know how I would parallelize this projection, and probably don't want to. But it has the advantage of being able to go between sine and cosine transforms. Any advice?

I'm not sure I can see the advantage without knowing more, but with mpi4py-fft you can configure any transform and the parallelization is taken care of for you:-)

francispoulin commented 4 years ago

Thanks. I did a bit of thinking about this and decided that maybe shenfun is the way to go, as you suggested a while ago. I'm going to try and build up from simple advection problems and see if I can get the hang of it to develop more complicated code. Cheers!