def circle_rand():
tmp_me = newDistArray(self.fft, self.params.spectral)
ndim = len(tmp_me.shape)
L = int(self.params.L)
# get random radii for circles/spheres
np.random.seed(1)
lbound = 3.0 * self.params.eps
ubound = 0.5 - self.params.eps
rand_radii = (ubound - lbound) * np.random.random_sample(
size=tuple([L] * ndim)) + lbound
# distribute circles/spheres
tmp = newDistArray(self.fft, False)
if ndim == 2:
for i in range(0, L):
for j in range(0, L):
# build radius
r2 = (self.X[0] + i - L + 0.5)**2 + (self.X[1] + j -
L + 0.5)**2
# add this blob, shifted by 1 to avoid issues with adding up negative contributions
tmp += np.tanh((rand_radii[i, j] - np.sqrt(r2)) /
(np.sqrt(2) * self.params.eps)) + 1
# normalize to [0,1]
tmp *= 0.5
assert np.all(tmp <= 1.0)
if self.params.spectral:
tmp_me[:] = self.fft.forward(tmp)
else:
tmp_me[:] = tmp[:]
return tmp_me
def eval_f(self, u, t):
"""
Routine to evaluate the RHS
Args:
u (dtype_u): current values
t (float): current time
Returns:
dtype_f: the RHS
"""
f = self.dtype_f(self.init)
if self.params.spectral:
f.impl[..., 0] = -self.K2 * u[..., 0]
if self.params.eps > 0:
tmp_u = newDistArray(self.fft, False)
tmp_T = newDistArray(self.fft, False)
tmp_u = self.fft.backward(u[..., 0], tmp_u)
tmp_T = self.fft.backward(u[..., 1], tmp_T)
tmpf = - 2.0 / self.params.eps ** 2 * tmp_u * (1.0 - tmp_u) * (1.0 - 2.0 * tmp_u) - \
6.0 * self.params.dw * (tmp_T - self.params.TM) / self.params.TM * tmp_u * (1.0 - tmp_u)
f.expl[..., 0] = self.fft.forward(tmpf)
f.impl[..., 1] = -self.params.D * self.K2 * u[..., 1]
f.expl[..., 1] = -f.impl[..., 0] - f.expl[..., 0]
else:
u_hat = self.fft.forward(u[..., 0])
lap_u_hat = -self.K2 * u_hat
f.impl[..., 0] = self.fft.backward(lap_u_hat, f.impl[..., 0])
if self.params.eps > 0:
f.expl[..., 0] = -2.0 / self.params.eps**2 * u[..., 0] * (
1.0 - u[..., 0]) * (1.0 - 2.0 * u[..., 0])
f.expl[..., 0] -= 6.0 * self.params.dw * (u[..., 1] - self.params.TM) / self.params.TM * u[..., 0] * \
(1.0 - u[..., 0])
u_hat = self.fft.forward(u[..., 1])
lap_u_hat = -self.params.D * self.K2 * u_hat
f.impl[..., 1] = self.fft.backward(lap_u_hat, f.impl[..., 1])
f.expl[..., 1] = -f.impl[..., 0] - f.expl[..., 0]
return f
def post_step(self, step, level_number):
"""
Overwrite standard post step hook
Args:
step (pySDC.Step.step): the current step
level_number (int): the current level number
"""
super(dump, self).post_step(step, level_number)
# some abbreviations
L = step.levels[0]
# get real space values
if L.prob.params.spectral:
if hasattr(L.prob, 'ncomp'):
tmp1 = newDistArray(L.prob.fft, False)
tmp = np.zeros(tmp1.shape + (L.prob.ncomp,))
for i in range(L.prob.ncomp):
tmp[..., i] = L.prob.fft.backward(L.uend[..., i])
else:
tmp = L.prob.fft.backward(L.uend)
else:
tmp = L.uend[:]
# compute local offset for I/O
nbytes_local = tmp.nbytes
if self.comm is not None:
nbytes_global = self.comm.allgather(nbytes_local)
else:
nbytes_global = [nbytes_local]
local_offset = sum(nbytes_global[:self.rank])
# dump initial data
fname = f"./data/{L.prob.params.name}_{self.time_step + step.status.slot:08d}"
fh = MPI.File.Open(self.comm, fname + ".dat", self.amode)
fh.Write_at_all(local_offset, tmp)
fh.Close()
sizes = list(L.prob.params.nvars)
if hasattr(L.prob, 'ncomp'):
sizes.append(L.prob.ncomp)
# write json description
if self.rank == 0:
json_obj = dict()
json_obj['type'] = 'dataset'
json_obj['datatype'] = str(tmp.dtype)
json_obj['endian'] = str(tmp.dtype.byteorder)
json_obj['time'] = L.time + L.dt
json_obj['space_comm_size'] = self.size
json_obj['time_comm_size'] = step.status.time_size
json_obj['shape'] = sizes
json_obj['elementsize'] = tmp.dtype.itemsize
with open(fname + '.json', 'w') as fp:
json.dump(json_obj, fp)
# update step count
self.time_step += step.status.time_size
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 restrict(self, F):
"""
Restriction implementation
Args:
F: the fine level data (easier to access than via the fine attribute)
"""
if isinstance(F, parallel_mesh):
if self.spectral:
G = self.coarse_prob.dtype_u(self.coarse_prob.init)
if hasattr(self.fine_prob, 'ncomp'):
for i in range(self.fine_prob.ncomp):
tmpF = newDistArray(self.fine_prob.fft, False)
tmpF = self.fine_prob.fft.backward(F[..., i], tmpF)
tmpG = tmpF[::int(self.ratio[0]), ::int(self.ratio[1])]
G[..., i] = self.coarse_prob.fft.forward(tmpG, G[..., i])
else:
tmpF = self.fine_prob.fft.backward(F)
tmpG = tmpF[::int(self.ratio[0]), ::int(self.ratio[1])]
G[:] = self.coarse_prob.fft.forward(tmpG, G)
else:
G = self.coarse_prob.dtype_u(self.coarse_prob.init)
G[:] = F[::int(self.ratio[0]), ::int(self.ratio[1])]
else:
raise TransferError('Unknown data type, got %s' % type(F))
return G
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 __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 u_exact(self, t):
"""
Routine to compute the exact solution at time t
Args:
t (float): current time
Returns:
dtype_u: exact solution
"""
assert t == 0, 'ERROR: u_exact only valid for t=0'
me = self.dtype_u(self.init, val=0.0)
if self.params.init_type == 'circle':
r2 = (self.X[0] - 0.5)**2 + (self.X[1] - 0.5)**2
if self.params.spectral:
tmp = 0.5 * (1.0 + np.tanh((self.params.radius - np.sqrt(r2)) /
(np.sqrt(2) * self.params.eps)))
me[:] = self.fft.forward(tmp)
else:
me[:] = 0.5 * (1.0 + np.tanh(
(self.params.radius - np.sqrt(r2)) /
(np.sqrt(2) * self.params.eps)))
elif self.params.init_type == 'circle_rand':
ndim = len(me.shape)
L = int(self.params.L)
# get random radii for circles/spheres
np.random.seed(1)
lbound = 3.0 * self.params.eps
ubound = 0.5 - self.params.eps
rand_radii = (ubound - lbound) * np.random.random_sample(
size=tuple([L] * ndim)) + lbound
# distribute circles/spheres
tmp = newDistArray(self.fft, False)
if ndim == 2:
for i in range(0, L):
for j in range(0, L):
# build radius
r2 = (self.X[0] + i - L + 0.5)**2 + (self.X[1] + j -
L + 0.5)**2
# add this blob, shifted by 1 to avoid issues with adding up negative contributions
tmp += np.tanh((rand_radii[i, j] - np.sqrt(r2)) /
(np.sqrt(2) * self.params.eps)) + 1
# normalize to [0,1]
tmp *= 0.5
assert np.all(tmp <= 1.0)
if self.params.spectral:
me[:] = self.fft.forward(tmp)
else:
me[:] = tmp[:]
else:
raise NotImplementedError(
'type of initial value not implemented, got %s' %
self.params.init_type)
return me
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 sine():
tmp_me = newDistArray(self.fft, self.params.spectral)
if self.params.spectral:
tmp = np.sin(2 * np.pi * self.X[0]) * np.sin(
2 * np.pi * self.X[1])
tmp_me[:] = self.fft.forward(tmp)
else:
tmp_me[:] = np.sin(2 * np.pi * self.X[0]) * np.sin(
2 * np.pi * self.X[1])
return tmp_me
def eval_f(self, u, t):
"""
Routine to evaluate the RHS
Args:
u (dtype_u): current values
t (float): current time
Returns:
dtype_f: the RHS
"""
f = self.dtype_f(self.init)
if self.params.spectral:
f.comp1[..., 0] = self.Ku * u[..., 0]
f.comp1[..., 1] = self.Kv * u[..., 1]
tmpu = newDistArray(self.fft, False)
tmpv = newDistArray(self.fft, False)
tmpu[:] = self.fft.backward(u[..., 0], tmpu)
tmpv[:] = self.fft.backward(u[..., 1], tmpv)
tmpfu = -tmpu * tmpv**2 + self.params.A * (1 - tmpu)
tmpfv = tmpu * tmpv**2 - self.params.B * tmpv
f.comp2[..., 0] = self.fft.forward(tmpfu)
f.comp2[..., 1] = self.fft.forward(tmpfv)
else:
u_hat = self.fft.forward(u[..., 0])
lap_u_hat = self.Ku * u_hat
f.comp1[..., 0] = self.fft.backward(lap_u_hat, f.comp1[..., 0])
u_hat = self.fft.forward(u[..., 1])
lap_u_hat = self.Kv * u_hat
f.comp1[..., 1] = self.fft.backward(lap_u_hat, f.comp1[..., 1])
f.comp2[...,
0] = -u[..., 0] * u[...,
1]**2 + self.params.A * (1 - u[..., 0])
f.comp2[...,
1] = u[..., 0] * u[..., 1]**2 - self.params.B * u[..., 1]
return f
def circle():
tmp_me = newDistArray(self.fft, self.params.spectral)
r2 = (self.X[0] - 0.5)**2 + (self.X[1] - 0.5)**2
if self.params.spectral:
tmp = 0.5 * (1.0 + np.tanh((self.params.radius - np.sqrt(r2)) /
(np.sqrt(2) * self.params.eps)))
tmp_me[:] = self.fft.forward(tmp)
else:
tmp_me[:] = 0.5 * (1.0 + np.tanh(
(self.params.radius - np.sqrt(r2)) /
(np.sqrt(2) * self.params.eps)))
return tmp_me
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)
curl_pad = newDistArray(FFT_pad, False, rank=1, view=True)
def get_local_mesh(FFT, L):
"""Returns local mesh."""
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)
u_hat = fft.forward(u, u_hat)
uj = np.zeros_like(u)
uj = fft.backward(u_hat, uj)
assert np.allclose(uj, u)
u_padded = newDistArray(pfft, forward_output=False)
uc = u_hat.copy()
u_padded = pfft.backward(u_hat, u_padded)
u_hat = pfft.forward(u_padded, u_hat)
assert np.allclose(u_hat, uc)
#cfft = PFFT(MPI.COMM_WORLD, N, dtype=complex, padding=[1.5, 1.5, 1.5])
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)
pf.backward(uk, u)
plt.figure()
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]
def eval_f(self, u, t):
"""
Routine to evaluate the RHS
Args:
u (dtype_u): current values
t (float): current time
Returns:
dtype_f: the RHS
"""
f = self.dtype_f(self.init)
if self.params.spectral:
f.impl = -self.K2 * u
tmp = newDistArray(self.fft, False)
tmp[:] = self.fft.backward(u, tmp)
if self.params.eps > 0:
tmpf = -2.0 / self.params.eps**2 * tmp * (1.0 - tmp) * (
1.0 - 2.0 * tmp)
else:
tmpf = self.dtype_f(self.init, val=0.0)
# build sum over RHS without driving force
Rt_local = float(np.sum(self.fft.backward(f.impl) + tmpf))
if self.params.comm is not None:
Rt_global = self.params.comm.allreduce(sendobj=Rt_local,
op=MPI.SUM)
else:
Rt_global = Rt_local
# build sum over driving force term
Ht_local = float(np.sum(6.0 * tmp * (1.0 - tmp)))
if self.params.comm is not None:
Ht_global = self.params.comm.allreduce(sendobj=Ht_local,
op=MPI.SUM)
else:
Ht_global = Rt_local
# add/substract time-dependent driving force
if Ht_global != 0.0:
dw = Rt_global / Ht_global
else:
dw = 0.0
tmpf -= 6.0 * dw * tmp * (1.0 - tmp)
f.expl[:] = self.fft.forward(tmpf)
else:
u_hat = self.fft.forward(u)
lap_u_hat = -self.K2 * u_hat
f.impl[:] = self.fft.backward(lap_u_hat, f.impl)
if self.params.eps > 0:
f.expl = -2.0 / self.params.eps**2 * u * (1.0 - u) * (1.0 -
2.0 * u)
# build sum over RHS without driving force
Rt_local = float(np.sum(f.impl + f.expl))
if self.params.comm is not None:
Rt_global = self.params.comm.allreduce(sendobj=Rt_local,
op=MPI.SUM)
else:
Rt_global = Rt_local
# build sum over driving force term
Ht_local = float(np.sum(6.0 * u * (1.0 - u)))
if self.params.comm is not None:
Ht_global = self.params.comm.allreduce(sendobj=Ht_local,
op=MPI.SUM)
else:
Ht_global = Rt_local
# add/substract time-dependent driving force
if Ht_global != 0.0:
dw = Rt_global / Ht_global
else:
dw = 0.0
f.expl -= 6.0 * dw * u * (1.0 - u)
return f
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)
u[:] = np.sin(2 * np.pi * X[0]) * np.sin(2 * np.pi * X[1])
print(u.shape, X[0].shape)
# exit()
uex[:] = -2.0 * (2.0 * np.pi)**2 * np.sin(2 * np.pi * X[0]) * np.sin(
2 * np.pi * X[1])
u_hat = fft.forward(u)
lap_u_hat = -K2 * u_hat
lap_u = np.zeros_like(u)
lap_u = fft.backward(lap_u_hat, lap_u)
def solve_system_2(self, rhs, factor, u0, t):
"""
Newton-Solver for the second component
Args:
rhs (dtype_f): right-hand side for the linear system
factor (float) : abbrev. for the node-to-node stepsize (or any other factor required)
u0 (dtype_u): initial guess for the iterative solver (not used here so far)
t (float): current time (e.g. for time-dependent BCs)
Returns:
dtype_u: solution as mesh
"""
u = self.dtype_u(u0)
if self.params.spectral:
tmpu = newDistArray(self.fft, False)
tmpv = newDistArray(self.fft, False)
tmpu[:] = self.fft.backward(u[..., 0], tmpu)
tmpv[:] = self.fft.backward(u[..., 1], tmpv)
tmprhsu = newDistArray(self.fft, False)
tmprhsv = newDistArray(self.fft, False)
tmprhsu[:] = self.fft.backward(rhs[..., 0], tmprhsu)
tmprhsv[:] = self.fft.backward(rhs[..., 1], tmprhsv)
else:
tmpu = u[..., 0]
tmpv = u[..., 1]
tmprhsu = rhs[..., 0]
tmprhsv = rhs[..., 1]
# start newton iteration
n = 0
res = 99
while n < self.params.newton_maxiter:
# print(n, res)
# form the function g with g(u) = 0
tmpgu = tmpu - tmprhsu - factor * (-tmpu * tmpv**2 + self.params.A)
tmpgv = tmpv - tmprhsv - factor * (tmpu * tmpv**2)
# if g is close to 0, then we are done
res = max(np.linalg.norm(tmpgu, np.inf),
np.linalg.norm(tmpgv, np.inf))
if res < self.params.newton_tol:
break
# assemble dg
dg00 = 1 - factor * (-tmpv**2)
dg01 = -factor * (-2 * tmpu * tmpv)
dg10 = -factor * (tmpv**2)
dg11 = 1 - factor * (2 * tmpu * tmpv)
# interleave and unravel to put into sparse matrix
dg00I = np.ravel(np.kron(dg00, np.array([1, 0])))
dg01I = np.ravel(np.kron(dg01, np.array([1, 0])))
dg10I = np.ravel(np.kron(dg10, np.array([1, 0])))
dg11I = np.ravel(np.kron(dg11, np.array([0, 1])))
# put into sparse matrix
dg = sp.diags(dg00I, offsets=0) + sp.diags(dg11I, offsets=0)
dg += sp.diags(dg01I, offsets=1, shape=dg.shape) + sp.diags(
dg10I, offsets=-1, shape=dg.shape)
# interleave g terms to apply inverse to it
g = np.kron(tmpgu.flatten(), np.array([1, 0])) + np.kron(
tmpgv.flatten(), np.array([0, 1]))
# invert dg matrix
b = sp.linalg.spsolve(dg, g)
# update real-space vectors
tmpu[:] -= b[::2].reshape(self.params.nvars)
tmpv[:] -= b[1::2].reshape(self.params.nvars)
# increase iteration count
n += 1
if np.isnan(res) and self.params.stop_at_nan:
raise ProblemError(
'Newton got nan after %i iterations, aborting...' % n)
elif np.isnan(res):
self.logger.warning('Newton got nan after %i iterations...' % n)
if n == self.params.newton_maxiter:
self.logger.warning(
'Newton did not converge after %i iterations, error is %s' %
(n, res))
# self.newton_ncalls += 1
# self.newton_itercount += n
me = self.dtype_u(self.init)
if self.params.spectral:
me[..., 0] = self.fft.forward(tmpu)
me[..., 1] = self.fft.forward(tmpv)
else:
me[..., 0] = tmpu
me[..., 1] = tmpv
return me
def run_simulation(name='',
spectral=None,
nprocs_time=None,
nprocs_space=None,
dt=None,
cwd='.'):
"""
A test program to do PFASST runs for the AC equation with temperature-based forcing
(slightly inefficient, but will run for a few seconds only)
Args:
name (str): name of the run, will be used to distinguish different setups
spectral (bool): run in real or spectral space
nprocs_time (int): number of processors in time
nprocs_space (int): number of processors in space (None if serial)
dt (float): time-step size
cwd (str): current working directory
"""
# set MPI communicator
comm = MPI.COMM_WORLD
world_rank = comm.Get_rank()
world_size = comm.Get_size()
# split world communicator to create space-communicators
if nprocs_space is not None:
color = int(world_rank / nprocs_space)
else:
color = int(world_rank / 1)
space_comm = comm.Split(color=color)
space_rank = space_comm.Get_rank()
space_size = space_comm.Get_size()
assert world_size == space_size, 'This script cannot run parallel-in-time with MPI, only spatial parallelism'
# initialize level parameters
level_params = dict()
level_params['restol'] = 1E-12
level_params['dt'] = dt
level_params['nsweeps'] = [1]
# initialize sweeper parameters
sweeper_params = dict()
sweeper_params['collocation_class'] = CollGaussRadau_Right
sweeper_params['num_nodes'] = [3]
sweeper_params['QI'] = [
'LU'
] # For the IMEX sweeper, the LU-trick can be activated for the implicit part
sweeper_params['initial_guess'] = 'spread'
# initialize problem parameters
problem_params = dict()
problem_params['L'] = 1.0
problem_params['nvars'] = [(128, 128), (32, 32)]
problem_params['eps'] = [0.04]
problem_params['radius'] = 0.25
problem_params['TM'] = 1.0
problem_params['D'] = 0.1
problem_params['dw'] = [21.0]
problem_params['comm'] = space_comm
problem_params['name'] = name
problem_params['init_type'] = 'circle'
problem_params['spectral'] = spectral
# initialize step parameters
step_params = dict()
step_params['maxiter'] = 50
# initialize controller parameters
controller_params = dict()
controller_params[
'logger_level'] = 30 if space_rank == 0 else 99 # set level depending on rank
controller_params['predict_type'] = 'pfasst_burnin'
# fill description dictionary for easy step instantiation
description = dict()
description['problem_params'] = problem_params # pass problem parameters
description['sweeper_class'] = imex_1st_order
description['sweeper_params'] = sweeper_params # pass sweeper parameters
description['level_params'] = level_params # pass level parameters
description['step_params'] = step_params # pass step parameters
description['space_transfer_class'] = fft_to_fft
description['problem_class'] = allencahn_temp_imex
# set time parameters
t0 = 0.0
Tend = 1 * 0.001
if space_rank == 0:
out = f'---------> Running {name} with spectral={spectral} and {space_size} process(es) in space...'
print(out)
# instantiate controller
controller = controller_nonMPI(num_procs=nprocs_time,
controller_params=controller_params,
description=description)
# get initial values on finest level
P = controller.MS[0].levels[0].prob
uinit = P.u_exact(t0)
# call main function to get things done...
uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)
if space_rank == 0:
# convert filtered statistics of iterations count, sorted by time
iter_counts = sort_stats(filter_stats(stats, type='niter'),
sortby='time')
niters = np.mean(np.array([item[1] for item in iter_counts]))
out = f'Mean number of iterations: {niters:.4f}'
print(out)
# get setup time
timing = sort_stats(filter_stats(stats, type='timing_setup'),
sortby='time')
out = f'Setup time: {timing[0][1]:.4f} sec.'
print(out)
# get running time
timing = sort_stats(filter_stats(stats, type='timing_run'),
sortby='time')
out = f'Time to solution: {timing[0][1]:.4f} sec.'
print(out)
refname = f'{cwd}/data/AC-reference-tempforce_00001000'
with open(f'{refname}.json', 'r') as fp:
obj = json.load(fp)
array = np.fromfile(f'{refname}.dat', dtype=obj['datatype'])
array = array.reshape(obj['shape'], order='C')
if spectral:
ureal = newDistArray(P.fft, False)
ureal = P.fft.backward(uend[..., 0], ureal)
Treal = newDistArray(P.fft, False)
Treal = P.fft.backward(uend[..., 1], Treal)
err = max(np.amax(abs(ureal - array[..., 0])),
np.amax(abs(Treal - array[..., 1])))
else:
err = abs(array - uend)
out = f'...Done <---------\n'
print(out)
return err
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()
