def test_newDistArray():
N = (8, 8, 8)
pfft = PFFT(MPI.COMM_WORLD, N)
for forward_output in (True, False):
for view in (True, False):
for rank in (0, 1, 2):
a = newDistArray(pfft,
forward_output=forward_output,
rank=rank,
view=view)
if view is False:
assert isinstance(a, DistArray)
assert a.rank == rank
if rank == 0:
qfft = PFFT(MPI.COMM_WORLD, darray=a)
elif rank == 1:
qfft = PFFT(MPI.COMM_WORLD, darray=a[0])
else:
qfft = PFFT(MPI.COMM_WORLD, darray=a[0, 0])
qfft.destroy()
else:
assert isinstance(a, np.ndarray)
assert a.base.rank == rank
pfft.destroy()
def __init__(self,
Nx,
Ny,
padx=1.5,
pady=1.5,
num_in=6,
num_out=2,
fmultin=mult_in,
fmultout=mult_out62,
comm=MPI.COMM_WORLD):
self.comm = comm
self.num_in = num_in
self.num_out = num_out
self.nar = max(num_in, num_out)
self.ffti = PFFT(comm,
shape=(self.num_in, Nx, Ny),
axes=(1, 2),
grid=[1, -1, 1],
padding=[1, 1.5, 1.5],
collapse=False)
self.ffto = PFFT(comm,
shape=(self.num_out, Nx, Ny),
axes=(1, 2),
grid=[1, -1, 1],
padding=[1, 1.5, 1.5],
collapse=False)
self.datk = newDistArray(self.ffti, forward_output=True)
self.dat = newDistArray(self.ffti, forward_output=False)
lkx = np.r_[0:int(Nx / 2), -int(Nx / 2):0]
lky = np.r_[0:int(Ny / 2 + 1)]
self.kx = DistArray((Nx, int(Ny / 2 + 1)),
subcomm=(1, 0),
dtype=float,
alignment=0)
self.ky = DistArray((Nx, int(Ny / 2 + 1)),
subcomm=(1, 0),
dtype=float,
alignment=0)
self.kx[:], self.ky[:] = np.meshgrid(lkx[self.kx.local_slice()[0]],
lky[self.ky.local_slice()[1]],
indexing='ij')
self.ksqr = self.kx**2 + self.ky**2
self.fmultin = fmultin
self.fmultout = fmultout
def test_2D(backend, forward_output):
if backend == 'netcdf4':
assert forward_output is False
T = PFFT(comm, (N[0], N[1]))
for i, domain in enumerate([
None, ((0, np.pi), (0, 2 * np.pi)),
(np.arange(N[0], dtype=float) * 1 * np.pi / N[0],
np.arange(N[1], dtype=float) * 2 * np.pi / N[1])
]):
for rank in range(3):
filename = "".join(
('test2D_{}{}{}'.format(ex[i == 0], ex[forward_output],
rank), ending[backend]))
if backend == 'netcdf4':
remove_if_exists(filename)
u = newDistArray(T,
forward_output=forward_output,
val=1,
rank=rank)
hfile = writer[backend](filename, domain=domain)
assert hfile.backend() == backend
hfile.write(0, {'u': [u]})
hfile.write(1, {'u': [u]})
u.write(hfile, 'u', 2)
if rank > 0:
hfile.write(0, {'u': [u]}, as_scalar=True)
hfile.write(1, {'u': [u]}, as_scalar=True)
u.write(hfile, 'u', 2, as_scalar=True)
u.write('t' + filename, 'u', 0)
u.write('t' + filename, 'u', 0, [slice(None), 3])
if not forward_output and backend == 'hdf5' and comm.Get_rank(
) == 0:
generate_xdmf(filename)
generate_xdmf(filename, order='visit')
u0 = newDistArray(T, forward_output=forward_output, rank=rank)
read = reader[backend](filename)
read.read(u0, 'u', step=0)
u0.read(filename, 'u', 2)
u0.read(read, 'u', 2)
assert np.allclose(u0, u)
if backend == 'netcdf4': # Test opening file in mode 'a' when not existing
remove_if_exists('nctesta.nc')
_ = NCFile('nctesta.nc', domain=domain, mode='a')
T.destroy()
def test_4D(backend, forward_output):
if backend == 'netcdf4':
assert forward_output is False
T = PFFT(comm, (N[0], N[1], N[2], N[3]))
d0 = ((0, np.pi), (0, 2 * np.pi), (0, 3 * np.pi), (0, 4 * np.pi))
d1 = (np.arange(N[0], dtype=float) * 1 * np.pi / N[0],
np.arange(N[1], dtype=float) * 2 * np.pi / N[1],
np.arange(N[2], dtype=float) * 3 * np.pi / N[2],
np.arange(N[3], dtype=float) * 4 * np.pi / N[3])
for i, domain in enumerate([None, d0, d1]):
for rank in range(3):
filename = "".join(
('h5test4_{}{}{}'.format(ex[i == 0], ex[forward_output],
rank), ending[backend]))
if backend == 'netcdf4':
remove_if_exists('uv' + filename)
u = newDistArray(T, forward_output=forward_output, rank=rank)
v = newDistArray(T, forward_output=forward_output, rank=rank)
h0file = writer[backend]('uv' + filename, domain=domain)
u[:] = np.random.random(u.shape)
v[:] = 2
for k in range(3):
h0file.write(
k, {
'u':
[u, (u, [slice(None), 4,
slice(None),
slice(None)])],
'v': [v,
(v, [slice(None), slice(None), 5, 6])]
})
if not forward_output and backend == 'hdf5' and comm.Get_rank(
) == 0:
generate_xdmf('uv' + filename)
u0 = newDistArray(T, forward_output=forward_output, rank=rank)
read = reader[backend]('uv' + filename)
read.read(u0, 'u', step=0)
assert np.allclose(u0, u)
read.read(u0, 'v', step=0)
assert np.allclose(u0, v)
T.destroy()
def __init__(self, decomp, queue, grid_shape, dtype, **kwargs):
self.decomp = decomp
self.grid_shape = grid_shape
self.proc_shape = decomp.proc_shape
self.dtype = np.dtype(dtype)
self.is_real = self.dtype.kind == "f"
from pystella.fourier import get_complex_dtype_with_matching_prec
self.cdtype = get_complex_dtype_with_matching_prec(self.dtype)
from pystella.fourier import get_real_dtype_with_matching_prec
self.rdtype = get_real_dtype_with_matching_prec(self.dtype)
if self.proc_shape[0] > 1 and self.proc_shape[1] == 1:
slab = True
else:
slab = False
from mpi4py_fft.pencil import Subcomm
default_kwargs = dict(
axes=([0], [1], [2]),
threads=16,
backend="fftw",
collapse=True,
)
default_kwargs.update(kwargs)
comm = decomp.comm if slab else Subcomm(decomp.comm, self.proc_shape)
from mpi4py_fft import PFFT
self.fft = PFFT(comm,
grid_shape,
dtype=dtype,
slab=slab,
**default_kwargs)
self.fx = self.fft.forward.input_array
self.fk = self.fft.forward.output_array
slc = self.fft.local_slice(True)
self.sub_k = get_sliced_momenta(grid_shape, self.dtype, slc, queue)
def __init__(self, fine_prob, coarse_prob, params):
"""
Initialization routine
Args:
fine_prob: fine problem
coarse_prob: coarse problem
params: parameters for the transfer operators
"""
# invoke super initialization
super(fft_to_fft, self).__init__(fine_prob, coarse_prob, params)
assert self.fine_prob.params.spectral == self.coarse_prob.params.spectral
self.spectral = self.fine_prob.params.spectral
Nf = list(self.fine_prob.fft.global_shape())
Nc = list(self.coarse_prob.fft.global_shape())
self.ratio = [int(nf / nc) for nf, nc in zip(Nf, Nc)]
axes = tuple(range(len(Nf)))
self.fft_pad = PFFT(self.coarse_prob.params.comm, Nc, padding=self.ratio, axes=axes,
dtype=self.coarse_prob.fft.dtype(False),
slab=True)
from mpi4py_fft import PFFT, newDistArray
# Set viscosity, end time and time step
nu = 0.000625
T = 0.1
dt = 0.01
# Set global size of the computational box
M = 6
N = [2**M, 2**M, 2**M]
L = np.array(
[2 * np.pi, 4 * np.pi, 4 * np.pi], dtype=float
) # Needs to be (2*int)*pi in all directions (periodic) because of initialization
# Create instance of PFFT to perform parallel FFT + an instance to do FFT with padding (3/2-rule)
FFT = PFFT(MPI.COMM_WORLD, N, collapse=False)
#FFT_pad = PFFT(MPI.COMM_WORLD, N, padding=[1.5, 1.5, 1.5])
FFT_pad = FFT
# Declare variables needed to solve Navier-Stokes
U = newDistArray(FFT, False, rank=1, view=True) # Velocity
U_hat = newDistArray(FFT, rank=1, view=True) # Velocity transformed
P = newDistArray(FFT, False, view=True) # Pressure (scalar)
P_hat = newDistArray(FFT, view=True) # Pressure transformed
U_hat0 = newDistArray(FFT, rank=1, view=True) # Runge-Kutta work array
U_hat1 = newDistArray(FFT, rank=1, view=True) # Runge-Kutta work array
a = [1. / 6., 1. / 3., 1. / 3., 1. / 6.] # Runge-Kutta parameter
b = [0.5, 0.5, 1.] # Runge-Kutta parameter
dU = newDistArray(FFT, rank=1, view=True) # Right hand side of ODEs
curl = newDistArray(FFT, False, rank=1, view=True)
U_pad = newDistArray(FFT_pad, False, rank=1, view=True)
import numpy as np
from mpi4py import MPI
from mpi4py_fft import PFFT, newDistArray
from mpi4py_fft.fftw import dctn, idctn
# Set global size of the computational box
N = np.array([18, 18, 18], dtype=int)
dct = functools.partial(dctn, type=3)
idct = functools.partial(idctn, type=3)
transforms = {(1, 2): (dct, idct)}
fft = PFFT(MPI.COMM_WORLD,
N,
axes=None,
collapse=True,
grid=(-1, ),
transforms=transforms)
pfft = PFFT(MPI.COMM_WORLD,
N,
axes=((0, ), (1, 2)),
grid=(-1, ),
padding=[1.5, 1.0, 1.0],
transforms=transforms)
assert fft.axes == pfft.axes
u = newDistArray(fft, forward_output=False)
u[:] = np.random.random(u.shape).astype(u.dtype)
u_hat = newDistArray(fft, forward_output=True)
@author: ogurcan
"""
from mpi4py import MPI
from mpi4py_fft import PFFT, newDistArray
import numpy as np
import matplotlib.pylab as plt
howmany = 6
Nx, Ny = 128, 128
padx, pady = 3 / 2, 3 / 2
Npx, Npy = int(128 * padx), int(128 * pady)
comm = MPI.COMM_WORLD
pf = PFFT(comm,
shape=(howmany, Nx, Ny),
axes=(1, 2),
grid=[1, -1, 1],
padding=[1, 1.5, 1.5],
collapse=False)
u = newDistArray(pf, forward_output=False)
uk = newDistArray(pf, forward_output=True)
n, x, y = np.meshgrid(np.arange(0, howmany),
np.linspace(-1, 1, Npx),
np.linspace(-1, 1, Npy),
indexing='ij')
nl, xl, yl = n[u.local_slice()], x[u.local_slice()], y[u.local_slice()]
u[:] = np.sin(4 * np.pi * (xl + 2 * yl)) * np.exp(-xl**2 / 2 / 0.04 -
yl**2 / 2 / 0.08) * (nl - 3)
u0 = u.copy()
pf.forward(u, uk)
def __init__(self,
problem_params,
dtype_u=parallel_mesh,
dtype_f=parallel_imex_mesh):
"""
Initialization routine
Args:
problem_params (dict): custom parameters for the example
dtype_u: fft data type (will be passed to parent class)
dtype_f: fft data type wuth implicit and explicit parts (will be passed to parent class)
"""
if 'L' not in problem_params:
problem_params['L'] = 1.0
if 'init_type' not in problem_params:
problem_params['init_type'] = 'circle'
if 'comm' not in problem_params:
problem_params['comm'] = None
if 'dw' not in problem_params:
problem_params['dw'] = 0.0
# these parameters will be used later, so assert their existence
essential_keys = ['nvars', 'eps', 'L', 'radius', 'dw', 'spectral']
for key in essential_keys:
if key not in problem_params:
msg = 'need %s to instantiate problem, only got %s' % (
key, str(problem_params.keys()))
raise ParameterError(msg)
if not (isinstance(problem_params['nvars'], tuple)
and len(problem_params['nvars']) > 1):
raise ProblemError('Need at least two dimensions')
# Creating FFT structure
ndim = len(problem_params['nvars'])
axes = tuple(range(ndim))
self.fft = PFFT(problem_params['comm'],
list(problem_params['nvars']),
axes=axes,
dtype=np.float,
collapse=True)
# get test data to figure out type and dimensions
tmp_u = newDistArray(self.fft, problem_params['spectral'])
# invoke super init, passing the communicator and the local dimensions as init
super(allencahn_imex,
self).__init__(init=(tmp_u.shape, problem_params['comm'],
tmp_u.dtype),
dtype_u=dtype_u,
dtype_f=dtype_f,
params=problem_params)
L = np.array([self.params.L] * ndim, dtype=float)
# get local mesh
X = np.ogrid[self.fft.local_slice(False)]
N = self.fft.global_shape()
for i in range(len(N)):
X[i] = (X[i] * L[i] / N[i])
self.X = [np.broadcast_to(x, self.fft.shape(False)) for x in X]
# get local wavenumbers and Laplace operator
s = self.fft.local_slice()
N = self.fft.global_shape()
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(ndim):
Ks[i] = (Ks[i] * Lp[i]).astype(float)
K = [np.broadcast_to(k, self.fft.shape(True)) for k in Ks]
K = np.array(K).astype(float)
self.K2 = np.sum(K * K, 0, dtype=float)
# Need this for diagnostics
self.dx = self.params.L / problem_params['nvars'][0]
self.dy = self.params.L / problem_params['nvars'][1]
Ks = np.meshgrid(*K, indexing='ij', sparse=True)
Lp = 2 * np.pi / L
for i in range(ndim):
Ks[i] = (Ks[i] * Lp[i]).astype(float)
return [np.broadcast_to(k, FFT.shape(True)) for k in Ks]
comm = MPI.COMM_WORLD
subcomm = comm.Split()
print(subcomm)
nvars = 8
ndim = 2
axes = tuple(range(ndim))
N = np.array([nvars] * ndim, dtype=int)
print(N, axes)
fft = PFFT(subcomm, N, axes=axes, dtype=np.float, slab=True)
# L = np.array([2*np.pi] * ndim, dtype=float)
L = np.array([1] * ndim, dtype=float)
print(fft.subcomm)
X = get_local_mesh(fft, L)
K = get_local_wavenumbermesh(fft, L)
K = np.array(K).astype(float)
K2 = np.sum(K * K, 0, dtype=float)
u = newDistArray(fft, False)
print(type(u))
print(u.subcomm)
uex = newDistArray(fft, False)
def test_3D(backend, forward_output):
if backend == 'netcdf4':
assert forward_output is False
T = PFFT(comm, (N[0], N[1], N[2]))
d0 = ((0, np.pi), (0, 2 * np.pi), (0, 3 * np.pi))
d1 = (np.arange(N[0], dtype=float) * 1 * np.pi / N[0],
np.arange(N[1], dtype=float) * 2 * np.pi / N[1],
np.arange(N[2], dtype=float) * 3 * np.pi / N[2])
for i, domain in enumerate([None, d0, d1]):
for rank in range(3):
filename = ''.join(
('test_{}{}{}'.format(ex[i == 0], ex[forward_output],
rank), ending[backend]))
if backend == 'netcdf4':
remove_if_exists('uv' + filename)
remove_if_exists('v' + filename)
u = newDistArray(T, forward_output=forward_output, rank=rank)
v = newDistArray(T, forward_output=forward_output, rank=rank)
h0file = writer[backend]('uv' + filename, domain=domain)
h1file = writer[backend]('v' + filename, domain=domain)
u[:] = np.random.random(u.shape)
v[:] = 2
for k in range(3):
h0file.write(
k, {
'u': [
u, (u, [slice(None), slice(None), 4]),
(u, [5, 5, slice(None)])
],
'v': [v,
(v, [slice(None), 6, slice(None)])]
})
h1file.write(
k, {
'v': [
v,
(v, [slice(None), 6, slice(None)]),
(v, [6, 6, slice(None)])
]
})
# One more time with same k
h0file.write(
k, {
'u': [
u, (u, [slice(None), slice(None), 4]),
(u, [5, 5, slice(None)])
],
'v': [v, (v, [slice(None), 6, slice(None)])]
})
h1file.write(
k, {
'v': [
v, (v, [slice(None), 6, slice(None)]),
(v, [6, 6, slice(None)])
]
})
if rank > 0:
for k in range(3):
u.write('uv' + filename, 'u', k, as_scalar=True)
u.write('uv' + filename,
'u',
k, [slice(None), slice(None), 4],
as_scalar=True)
u.write('uv' + filename,
'u',
k, [5, 5, slice(None)],
as_scalar=True)
v.write('uv' + filename, 'v', k, as_scalar=True)
v.write('uv' + filename,
'v',
k, [slice(None), 6, slice(None)],
as_scalar=True)
if not forward_output and backend == 'hdf5' and comm.Get_rank(
) == 0:
generate_xdmf('uv' + filename)
generate_xdmf('v' + filename, periodic=False)
generate_xdmf('v' + filename, periodic=(True, True, True))
generate_xdmf('v' + filename, order='visit')
u0 = newDistArray(T, forward_output=forward_output, rank=rank)
read = reader[backend]('uv' + filename)
read.read(u0, 'u', step=0)
assert np.allclose(u0, u)
read.read(u0, 'v', step=0)
assert np.allclose(u0, v)
T.destroy()
o = mytype(m) m[0, 0] = 2 assert o[0, 0] == -1 assert o is not m exit() print(type(m)) print(type(m+n)) print(abs(m)) print(type(abs(m))) comm = MPI.COMM_WORLD subcomm = comm.Split() # print(subcomm) nvars = 8 ndim = 2 axes = tuple(range(ndim)) N = np.array([nvars] * ndim, dtype=int) # print(N, axes) fft = PFFT(subcomm, N, axes=axes, dtype=np.float, slab=True) init = (fft, False) m = fft_datatype(init) m[:] = comm.Get_rank() print(type(m)) print(m.subcomm) print(abs(m), type(abs(m)))
