def test_is_n_byte_aligned(self):
a = n_byte_align_empty(100, 16)
self.assertTrue(is_n_byte_aligned(a, 16))
a = n_byte_align_empty(100, 5)
self.assertTrue(is_n_byte_aligned(a, 5))
a = n_byte_align_empty(100, 16, dtype="float32")[1:]
self.assertFalse(is_n_byte_aligned(a, 16))
self.assertTrue(is_n_byte_aligned(a, 4))
def test_is_n_byte_aligned(self):
a = n_byte_align_empty(100, 16)
self.assertTrue(is_n_byte_aligned(a, 16))
a = n_byte_align_empty(100, 5)
self.assertTrue(is_n_byte_aligned(a, 5))
a = n_byte_align_empty(100, 16, dtype='float32')[1:]
self.assertFalse(is_n_byte_aligned(a, 16))
self.assertTrue(is_n_byte_aligned(a, 4))
def fftw_inplace(E):
if np.any(np.array(E.shape[2:]) > 1):
strides_not_identical = np.any(E.strides != fftw_vec_array.strides)
if strides_not_identical:
log.debug('In-place Fourier Transform: strides not identical.')
E = self.__word_align(E.copy())
if not pyfftw.is_n_byte_aligned(E, pyfftw.simd_alignment):
log.debug('In-place Fourier Transform: Input/Output array not %d-byte word aligned, aligning now.'
% pyfftw.simd_alignment)
E = self.__word_align(E)
assert(E.ndim-2 == self.data_shape.size and np.all(E.shape[2:] == self.data_shape))
if np.all(E.shape == fft_vec_object.input_shape) \
and pyfftw.is_n_byte_aligned(E, pyfftw.simd_alignment):
fft_vec_object(E, E) # E should be in SIMD-word-aligned memory zone
else:
log.debug('Fourier Transform: Array shape not standard, falling back to default interface.')
E = ft.fftn(E, axes=ft_axes)
return E
def check_alignment(array, message='NOT ALIGNED'):
"""
Diagnostic method for testing the word-alignment of an array to enable efficient operations.
:param array: The ndarray to be tested.
:param message: Optional message to display in case of failure.
"""
aligned = pyfftw.is_n_byte_aligned(array, pyfftw.simd_alignment)
if not aligned and message is not None:
log.debug(message)
return aligned
def __init__(self,
shape,
dtype_in=None,
data_in=None,
overwrite=True,
inverse=True,
packed=True,
use_pyfftw=True):
try:
nx, ny, nz = shape
except (TypeError, ValueError):
raise ValueError('Expected 3D shape.')
if nx % 2 or ny % 2 or nz % 2:
raise ValueError('All shape dimensions must be even.')
if data_in is not None:
if not isinstance(data_in, np.ndarray):
raise ValueError('Invalid type for data_in: {0}.'.format(
type(data_in)))
dtype_in = data_in.dtype
# Convert dtype_in to an object in the numpy scalar type hierarchy.
dtype_in = np.obj2sctype(dtype_in)
if dtype_in is None:
raise ValueError('Invalid dtype_in: {0}.'.format(dtype_in))
# Determine the input and output array type and shape.
if packed:
if inverse:
shape_in = (nx, ny, nz // 2 + 1)
if not issubclass(dtype_in, np.complexfloating):
raise ValueError(
'Invalid dtype_in for inverse packed transform ' +
'(should be complex): {0}.'.format(dtype_in))
dtype_out = scalar_type(dtype_in)
shape_out = (nx, ny, nz + 2) if overwrite else shape
else:
shape_in = (nx, ny, nz + 2) if overwrite else shape
if not issubclass(dtype_in, np.floating):
raise ValueError(
'Invalid dtype_in for forward packed transform ' +
'(should be floating): {0}.'.format(dtype_in))
dtype_out = complex_type(dtype_in)
shape_out = (nx, ny, nz // 2 + 1)
else:
if not issubclass(dtype_in, np.complexfloating):
raise ValueError(
'Expected complex dtype_in for transform: {0}.'.format(
dtype_in))
shape_in = shape_out = shape
dtype_out = dtype_in
if data_in is not None:
if data_in.shape != shape_in:
raise ValueError(
'data_in has wrong shape {0}, expected {1}.'.format(
data_in.shape, shape_in))
self.data_in = data_in
self.nbytes_allocated = 0
else:
# Allocate the input and output data buffers.
self.data_in = allocate(shape_in, dtype_in, use_pyfftw=use_pyfftw)
self.nbytes_allocated = self.data_in.nbytes
if overwrite:
if packed:
# See https://github.com/hgomersall/pyFFTW/issues/29
self.data_out = self.data_in.view(dtype_out).reshape(shape_out)
# Hide the padding without copying. See http://www.fftw.org/doc/
# Multi_002dDimensional-DFTs-of-Real-Data.html.
if inverse:
self.data_out_padded = self.data_out
self.data_out = self.data_out[:, :, :nz]
else:
self.data_in_padded = self.data_in
self.data_in = self.data_in[:, :, :nz]
else:
self.data_out = self.data_in
else:
self.data_out = allocate(shape_out,
dtype_out,
use_pyfftw=use_pyfftw)
self.nbytes_allocated += self.data_out.nbytes
# Try to use pyFFTW to configure the transform, if requested.
self.use_pyfftw = use_pyfftw
if self.use_pyfftw:
try:
import pyfftw
if not pyfftw.is_n_byte_aligned(self.data_in,
pyfftw.simd_alignment):
raise ValueError('data_in is not SIMD aligned.')
if not pyfftw.is_n_byte_aligned(self.data_out,
pyfftw.simd_alignment):
raise ValueError('data_out is not SIMD aligned.')
direction = 'FFTW_BACKWARD' if inverse else 'FFTW_FORWARD'
self.fftw_plan = pyfftw.FFTW(self.data_in,
self.data_out,
direction=direction,
flags=('FFTW_ESTIMATE', ),
axes=(0, 1, 2))
self.fftw_norm = np.float(nx * ny * nz if inverse else 1)
except ImportError:
self.use_pyfftw = False
# Fall back to numpy.fft if we are not using pyFFTW.
if not self.use_pyfftw:
if inverse:
self.transformer = np.fft.irfftn if packed else np.fft.ifftn
else:
self.transformer = np.fft.rfftn if packed else np.fft.fftn
# Remember our options so we can create a reverse plan.
self.shape = shape
self.inverse = inverse
self.packed = packed
self.overwrite = overwrite
def _Xfftn(a, s, axes, overwrite_input,
planner_effort, threads, auto_align_input, auto_contiguous,
avoid_copy, inverse, real):
'''Generic transform interface for all the transforms. No
defaults exist. The transform must be specified exactly.
'''
a_orig = a
invreal = inverse and real
if inverse:
direction = 'FFTW_BACKWARD'
else:
direction = 'FFTW_FORWARD'
if planner_effort not in _valid_efforts:
raise ValueError('Invalid planner effort: ', planner_effort)
s, axes = _cook_nd_args(a, s, axes, invreal)
input_shape, output_shape = _compute_array_shapes(
a, s, axes, inverse, real)
a_is_complex = numpy.iscomplexobj(a)
# Make the input dtype correct
if a.dtype not in _rc_dtype_pairs:
# We make it the default dtype
if not real or inverse:
# It's going to be complex
a = numpy.asarray(a, dtype=_rc_dtype_pairs[_default_dtype])
else:
a = numpy.asarray(a, dtype=_default_dtype)
elif not (real and not inverse) and not a_is_complex:
# We need to make it a complex dtype
a = numpy.asarray(a, dtype=_rc_dtype_pairs[a.dtype])
elif (real and not inverse) and a_is_complex:
# It should be real
a = numpy.asarray(a, dtype=_rc_dtype_pairs[a.dtype])
# Make the output dtype correct
if not real:
output_dtype = a.dtype
else:
output_dtype = _rc_dtype_pairs[a.dtype]
if not avoid_copy:
a_copy = a.copy()
output_array = pyfftw.n_byte_align_empty(output_shape,
pyfftw.simd_alignment, output_dtype)
flags = [planner_effort]
if not auto_align_input:
flags.append('FFTW_UNALIGNED')
if overwrite_input:
flags.append('FFTW_DESTROY_INPUT')
if not a.shape == input_shape:
if avoid_copy:
raise ValueError('Cannot avoid copy: '
'The transform shape is not the same as the array size. '
'(from avoid_copy flag)')
# This means we need to use an _FFTWWrapper object
# and so need to create slicers.
update_input_array_slicer, FFTW_array_slicer = (
_setup_input_slicers(a.shape, input_shape))
# Also, the input array will be a different shape to the shape of
# `a`, so we need to create a new array.
input_array = pyfftw.n_byte_align_empty(input_shape,
pyfftw.simd_alignment, a.dtype)
FFTW_object = _FFTWWrapper(input_array, output_array, axes, direction,
flags, threads, input_array_slicer=update_input_array_slicer,
FFTW_array_slicer=FFTW_array_slicer)
# We copy the data back into the internal FFTW object array
internal_array = FFTW_object.input_array
internal_array[:] = 0
internal_array[FFTW_array_slicer] = (
a_copy[update_input_array_slicer])
else:
# Otherwise we can use `a` as-is
input_array = a
if auto_contiguous:
# We only need to create a new array if it's not already
# contiguous
if not (a.flags['C_CONTIGUOUS'] or a.flags['F_CONTIGUOUS']):
if avoid_copy:
raise ValueError('Cannot avoid copy: '
'The input array is not contiguous and '
'auto_contiguous is set. (from avoid_copy flag)')
input_array = pyfftw.n_byte_align_empty(a.shape,
pyfftw.simd_alignment, a.dtype)
if (auto_align_input and
not pyfftw.is_n_byte_aligned(input_array,
pyfftw.simd_alignment)):
if avoid_copy:
raise ValueError('Cannot avoid copy: '
'The input array is not aligned and '
'auto_align is set. (from avoid_copy flag)')
input_array = pyfftw.n_byte_align(input_array,
pyfftw.simd_alignment)
FFTW_object = pyfftw.FFTW(input_array, output_array, axes, direction,
flags, threads)
if not avoid_copy:
# Copy the data back into the (likely) destroyed array
FFTW_object.input_array[:] = a_copy
return FFTW_object
def test_differing_aligned_arrays_update(self):
'''Test to see if the alignment code is working as expected
'''
# Start by creating arrays that are only on various byte
# alignments (4, 16 and 32)
_input_array = n_byte_align_empty(
len(self.input_array.ravel())*2+5,
32, dtype='float32')
_output_array = n_byte_align_empty(
len(self.output_array.ravel())*2+5,
32, dtype='float32')
_input_array[:] = 0
_output_array[:] = 0
input_array_4 = (
numpy.frombuffer(_input_array[1:-4].data, dtype='complex64')
.reshape(self.input_array.shape))
output_array_4 = (
numpy.frombuffer(_output_array[1:-4].data, dtype='complex64')
.reshape(self.output_array.shape))
input_array_16 = (
numpy.frombuffer(_input_array[4:-1].data, dtype='complex64')
.reshape(self.input_array.shape))
output_array_16 = (
numpy.frombuffer(_output_array[4:-1].data, dtype='complex64')
.reshape(self.output_array.shape))
input_array_32 = (
numpy.frombuffer(_input_array[:-5].data, dtype='complex64')
.reshape(self.input_array.shape))
output_array_32 = (
numpy.frombuffer(_output_array[:-5].data, dtype='complex64')
.reshape(self.output_array.shape))
input_arrays = {4: input_array_4,
16: input_array_16,
32: input_array_32}
output_arrays = {4: output_array_4,
16: output_array_16,
32: output_array_32}
alignments = (4, 16, 32)
# Test the arrays are aligned on 4 bytes...
self.assertTrue(is_n_byte_aligned(input_arrays[4], 4))
self.assertTrue(is_n_byte_aligned(output_arrays[4], 4))
# ...and on 16...
self.assertFalse(is_n_byte_aligned(input_arrays[4], 16))
self.assertFalse(is_n_byte_aligned(output_arrays[4], 16))
self.assertTrue(is_n_byte_aligned(input_arrays[16], 16))
self.assertTrue(is_n_byte_aligned(output_arrays[16], 16))
# ...and on 32...
self.assertFalse(is_n_byte_aligned(input_arrays[16], 32))
self.assertFalse(is_n_byte_aligned(output_arrays[16], 32))
self.assertTrue(is_n_byte_aligned(input_arrays[32], 32))
self.assertTrue(is_n_byte_aligned(output_arrays[32], 32))
if len(pyfftw.pyfftw._valid_simd_alignments) > 0:
max_align = pyfftw.pyfftw._valid_simd_alignments[0]
else:
max_align = simd_alignment
for in_align in alignments:
for out_align in alignments:
expected_align = min(in_align, out_align, max_align)
fft = FFTW(input_arrays[in_align], output_arrays[out_align])
self.assertTrue(fft.input_alignment == expected_align)
self.assertTrue(fft.output_alignment == expected_align)
for update_align in alignments:
if update_align < expected_align:
# This should fail (not aligned properly)
self.assertRaisesRegex(ValueError,
'Invalid input alignment',
fft.update_arrays,
input_arrays[update_align],
output_arrays[out_align])
self.assertRaisesRegex(ValueError,
'Invalid output alignment',
fft.update_arrays,
input_arrays[in_align],
output_arrays[update_align])
else:
# This should work (and not segfault!)
fft.update_arrays(input_arrays[update_align],
output_arrays[out_align])
fft.update_arrays(input_arrays[in_align],
output_arrays[update_align])
fft.execute()
N_THREADS = 2
print("N threads: %d" % N_THREADS)
al = pyfftw.simd_alignment
psi = pyfftw.n_byte_align_empty((N, N), al, 'complex128')
flag = 'FFTW_PATIENT'
fft_object = pyfftw.FFTW(psi, psi, flags=[flag], axes=(0, 1), threads=N_THREADS)
ifft_object = pyfftw.FFTW(psi, psi, flags=[flag], axes=(0, 1),
threads=N_THREADS, direction='FFTW_BACKWARD')
# copy psi0 into psi. To be safe about keeping the alignment of psi, set all the
# entries to 1 and then multiply.
psi[:] = 1.0
psi *= psi0
# Check psi is aligned
assert pyfftw.is_n_byte_aligned(psi, al)
print("Planning finished")
# Run simulation
for step in np.arange(N_TIMESTEPS):
# Implementing split-step method
# Update wavefunction and resovoir, record density
# Take fft, multiply by kinetic factor and then take inverse fft.
fft_object()
psi *= kineticFactorHalf
ifft_object()
# psi = fft.ifft2(kineticFactorHalf * fft.fft2(psi))
currentDensity = np.absolute(psi) ** 2
expFactorExciton = np.exp(- (gamma_R + R * currentDensity) * dt)
n *= expFactorExciton
def solve(self, params, pumpFunction, psi0Function=None,
potentialFunction=None, dt=0.1, T_MAX=50, continuing=False,
method='split-step', recordEnergy=True, stepsPerObservation=10):
"""
Solve the GPE for the given conditions. Records the final condensate
wavefunction and density in the GPESolver object.
Parameters:
params: A dictionary containing at leas the parameters of the
normalised GPE equation. That is, g_C, g_R, gamma_C, gamma_R, R, and
Pth
dt: the time increment per timestep. In units of the characteristic
time defined in ParameterContainer
T_MAX: the maximum time. In units of the characteristic time defined
in ParameterContainer.
pumpFunction: A function that describes the pump power in units of
the threshold power.
potentialFunction: A function that describes the potential in units
of the characteristic energy scale. Default is 0
psi0Function: A function that describes the initial wavefunction in
units of inverse characteristic length. Default is a random seed
whose real and imaginary parts are both drawn from a normal
distributions with a small (10e-2) mean and (10e-3) standard
deviation.
continue: Flag used to continue a simulation. If true, the inital
self.psiFinal will be taken as the initial wavefunction, and the
simulation will be continued in steps of dt until T_MAX is reached
"""
# TODO Add support for making videos, recording the energy, etc. as
# the need arises
# Check we have the required parameters and assign them
for key in self.__class__.__requiredParams:
if key not in params:
raise ValueError("Required Parameter %s missing" % key)
self.__setattr__(key, params[key])
# Check that we are not continuing and providing psi0
if continuing and psi0Function:
raise ValueError("Can't continue a simulation and provide psi0")
# Assign the potential, pump, etc.
P = pumpFunction(self.x, self.y)
if not continuing:
self.time = 0.0
if not psi0Function:
psi0 = (np.abs(np.random.normal(size=(self.N, self.N),
scale=10e-5, loc=10e-4))
+ 0.0j*np.random.normal(size=(self.N, self.N),
scale=10e-4, loc=10e-3))
else:
psi0 = psi0Function(self.x, self.y)
else:
psi0 = np.copy(self.psiFinal)
if not potentialFunction:
potential = 0.0
else:
potential = potentialFunction(self.x, self.y)
# Carefully copy psi0 into psi. We want to be certain that we keep the
# alignment
self.psi[:] = 1
self.psi *= psi0
# Check psi is aligned
assert pyfftw.is_n_byte_aligned(self.psi, self.al)
n = np.zeros_like(self.psi, dtype=np.float64) + 1.0
# Get kinetic factor, initial density
# TODO: Is the factor of 1/2 correct?
kineticFactorHalf = np.exp(-1.0j * self.k * self.K * dt / 2.0)
currentDensity = np.absolute(psi0) ** 2
# Do the simulation
N_TIMESTEPS = int((T_MAX - self.time) // dt)
# Set up array to record energy
if recordEnergy:
self.energy = np.zeros(N_TIMESTEPS // stepsPerObservation + 1)
self.energyTimes = np.zeros_like(self.energy)
if method == 'split-step':
for step in xrange(N_TIMESTEPS):
# Implementing split-step method
# Record energy
if step % stepsPerObservation == 0 and recordEnergy:
density = np.absolute(self.psi) ** 2
gradx = np.gradient(self.psi)[0]
normFactor = density.sum()
self.energyTimes[step // stepsPerObservation] = self.time
self.energy[step // stepsPerObservation] = \
-(0.25 * np.gradient(
np.gradient(density)[0])[0]
- 0.5 * np.absolute(gradx) ** 2
- (self.g_C * density + self.g_R * n)
* density).sum() / normFactor
# Update wavefunction and resovoir, record density
# Take fft, multiply by kinetic factor and then take inverse fft
self.fft_object()
self.psi *= kineticFactorHalf
self.ifft_object()
currentDensity = np.absolute(self.psi) ** 2
expFactorExciton = np.exp(- (self.gamma_R
+ self.R * currentDensity) * dt)
n *= expFactorExciton
n += P * dt
expFactorPolariton = np.exp((n*(0.5 * self.R - 1.0j * self.g_R)
- 0.5 * self.gamma_C
- 1.0j * self.g_C * currentDensity
- 1.0j * potential) * dt)
# Do the nonlinear update, take FFT, do kinetic update, and then
# take the inverse fft
self.psi *= expFactorPolariton
self.fft_object()
self.psi *= kineticFactorHalf
self.ifft_object()
self.psi *= self.damping
n *= self.damping
self.time += dt
if step % 100 == 0:
print("Time = %f" % (self.time))
elif method == 'RK4':
def nonLinearUpdateN(n, P, density, dt):
"""
Updates n (in place) to M(n, P, density, dt)
"""
n *= - (self.gamma_R + self.R * density) * dt
n += P * dt
def nonLinearUpdatePsi(psi, potential, density, n, dt):
psi *= -1j * dt * (n * (0.5j * self.R + self.g_R) -
0.5j * self.gamma_C
+ self.g_C * density + potential)
def diffusionUpdate(psi, psifft, psiifft, expFactor):
psifft()
psi *= expFactor
psiifft()
# Set up psiK nK, n0
kineticFactorRK4 = np.exp(-1.0j * self.K * dt / 2.0)
psiK = pyfftw.n_byte_align_empty((self.N, self.N), self.al,
'complex128')
n0 = np.zeros_like(n)
nK = np.zeros_like(n)
nkStored = np.zeros_like(n)
psiI = np.zeros_like(self.psi)
dK = np.absolute(self.psi) ** 2
# Set up an fft object for psi_K
fft_objectpsiK = pyfftw.FFTW(psiK, psiK,
flags=self.fft_object.flags,
axes=(0, 1), threads=self.N_THREADS)
ifft_objectpsiK = pyfftw.FFTW(psiK, psiK,
flags=self.fft_object.flags,
axes=(0, 1), threads=self.N_THREADS,
direction="FFTW_BACKWARD")
psiK[:] = self.psi
assert pyfftw.is_n_byte_aligned(self.psi, self.al)
assert pyfftw.is_n_byte_aligned(psiK, self.al)
for step in xrange(N_TIMESTEPS):
n0[:] = n
nK[:] = n
psiK[:] = self.psi
dK[:] = np.absolute(self.psi) ** 2
assert pyfftw.is_n_byte_aligned(psiK, self.al)
diffusionUpdate(self.psi, self.fft_object, self.ifft_object,
kineticFactorRK4)
psiI[:] = self.psi
# K=1 step
# Updates on psiK and nK
nonLinearUpdatePsi(psiK, potential, dK,
nK, dt)
nonLinearUpdateN(nK, P, dK, dt)
diffusionUpdate(psiK, fft_objectpsiK, ifft_objectpsiK,
kineticFactorRK4)
# Update the next value of psi and n
self.psi += psiK / 6
n += nK / 6
# Update psik and nK. Store things for the next step
psiK /= 2
psiK += psiI
nK /= 2
nK += n0
dK[:] = np.absolute(psiK) ** 2
self.time += dt / 2
# K=2 step
# Do updates
nonLinearUpdatePsi(psiK, potential, dK, nK, dt)
nonLinearUpdateN(nK, P, dK, dt)
# Update the next value of psi and n
self.psi += psiK / 3
n += nK / 3
# Update psiK and nK. Store things for the next step
psiK /= 2
psiK += psiI
nK /= 2
nK += n0
dK[:] = np.absolute(psiK) ** 2
# K=3 step
# Do the updates
nonLinearUpdatePsi(psiK, potential, dK,
nkStored, dt)
nonLinearUpdateN(nK, P, dK, dt)
# Update the next value of psi and n
self.psi += psiK / 3
n += nK / 3
# Update psiK and nK. Store things for the next step
psiK += psiI
nK += n0
dK[:] = np.absolute(psiK) ** 2
self.time += dt / 2
# K=4 step
# Do the updates
diffusionUpdate(psiK, fft_objectpsiK, ifft_objectpsiK,
kineticFactorRK4)
nonLinearUpdatePsi(psiK, potential, dK, nK, dt)
nonLinearUpdateN(nK, P, dK, dt)
# Transform psi back
diffusionUpdate(self.psi, self.fft_object, self.ifft_object,
kineticFactorRK4)
# Update the next value of psi and n
self.psi += psiK / 6
n += nK / 6
# Apply damping
self.psi *= self.damping
n *= self.damping
print("Time: %f" % self.time)
elif method == "RK4Exact":
def updateNtoF(n, density, dt):
"""
Updates n to F.
"""
# We have to do the assignment like this in order to actually
# change n
n[:] = np.exp(- (self.gamma_R + self.R * density) * dt)
def nonLinearUpdateF(F, n0, P, density, dt):
"""
Updates F (in place) to F **(0.5) * n0 + P*dt/2
"""
F **= 0.5
F *= n0
F += 0.5 * P * dt
def nonLinearUpdatePsi(psi, potential, density, n, dt):
psi *= -1j * dt * (n * (0.5j * self.R + self.g_R) -
0.5j * self.gamma_C
+ self.g_C * density + potential)
def diffusionUpdate(psi, psifft, psiifft, expFactor):
psifft()
psi *= expFactor
psiifft()
# Set up psiK nK, n0
kineticFactorRK4 = np.exp(-1.0j * self.K * dt / 2.0)
psiK = pyfftw.n_byte_align_empty((self.N, self.N), self.al,
'complex128')
n0 = np.zeros_like(n)
nK = np.zeros_like(n)
psiI = np.zeros_like(self.psi)
dK = np.absolute(self.psi) ** 2
# Set up an fft object for psi_K
fft_objectpsiK = pyfftw.FFTW(psiK, psiK,
flags=self.fft_object.flags,
axes=(0, 1), threads=self.N_THREADS)
ifft_objectpsiK = pyfftw.FFTW(psiK, psiK,
flags=self.fft_object.flags,
axes=(0, 1), threads=self.N_THREADS,
direction="FFTW_BACKWARD")
psiK[:] = self.psi
assert pyfftw.is_n_byte_aligned(self.psi, self.al)
assert pyfftw.is_n_byte_aligned(psiK, self.al)
for step in xrange(N_TIMESTEPS):
n0[:] = n
nK[:] = n
#
n[:] = 0
psiK[:] = self.psi
dK[:] = np.absolute(self.psi) ** 2
assert pyfftw.is_n_byte_aligned(psiK, self.al)
diffusionUpdate(self.psi, self.fft_object, self.ifft_object,
kineticFactorRK4)
psiI[:] = self.psi
# K=1 step
# Updates on psiK and nK
nonLinearUpdatePsi(psiK, potential, dK,
nK, dt)
diffusionUpdate(psiK, fft_objectpsiK, ifft_objectpsiK,
kineticFactorRK4)
updateNtoF(nK, dK, dt)
# Update the next value of psi and n
self.psi += psiK / 6
n += (nK*n0 + P*dt) / 6
# Update psik and nK. Store things for the next step
psiK /= 2
psiK += psiI
nonLinearUpdateF(nK, n0, P, dK, dt)
dK[:] = np.absolute(psiK) ** 2
self.time += dt / 2
# K=2 step
# Do updates
nonLinearUpdatePsi(psiK, potential, dK, nK, dt)
updateNtoF(nK, dK, dt)
# Update the next value of psi and n
self.psi += psiK / 3
n += (nK*n0 + P*dt) / 3
# Update psiK and nK. Store things for the next step
psiK /= 2
psiK += psiI
nonLinearUpdateF(nK, n0, P, dK, dt)
dK[:] = np.absolute(psiK) ** 2
# K=3 step
# Do the updates
nonLinearUpdatePsi(psiK, potential, dK,
nK, dt)
updateNtoF(nK, dK, dt)
# For the next step, the derivatives will be evaluated at the
# end of the timestep. So we need to update nK accordingly. This
# value of nK is also the one to use to update n
nK *= n0
nK += P*dt
# Update the next value of psi and n
self.psi += psiK / 3
n += nK / 3
# Update psiK and nK. Store things for the next step
# No need to update nK. Explained above
psiK += psiI
dK[:] = np.absolute(psiK) ** 2
self.time += dt / 2
# K=4 step
# Do the updates
diffusionUpdate(psiK, fft_objectpsiK, ifft_objectpsiK,
kineticFactorRK4)
nonLinearUpdatePsi(psiK, potential, dK, nK, dt)
updateNtoF(nK, dK, dt)
# Can update nK to the value at the end of the step as we don't
# need to store the half-step estimate since we don't have to do
# any more steps after this.
nK *= n0
nK += P*dt
# Transform psi back
diffusionUpdate(self.psi, self.fft_object, self.ifft_object,
kineticFactorRK4)
# Update the next value of psi and n
self.psi += psiK / 6
n += nK / 6
# Apply damping
self.psi *= self.damping
n *= self.damping
print("Time: %f" % self.time)
elif method == 'RK4LessExact':
def nonLinearUpdateN(n, P, density, dt):
"""
Updates n (in place) to M(n, P, density, dt)
"""
n *= np.exp(- (self.gamma_R + self.R * density) * dt)
n += P * dt
def nonLinearUpdatePsi(psi, potential, density, n, dt):
psi *= -1j * dt * (n * (0.5j * self.R + self.g_R) -
0.5j * self.gamma_C
+ self.g_C * density + potential)
def diffusionUpdate(psi, psifft, psiifft, expFactor):
psifft()
psi *= expFactor
psiifft()
# Set up psiK nK, n0
kineticFactorRK4 = np.exp(-1.0j * self.K * dt / 2.0)
psiK = pyfftw.n_byte_align_empty((self.N, self.N), self.al,
'complex128')
dK = np.zeros_like(n)
psiI = np.zeros_like(self.psi)
# Set up an fft object for psi_K
fft_objectpsiK = pyfftw.FFTW(psiK, psiK,
flags=self.fft_object.flags,
axes=(0, 1), threads=self.N_THREADS)
ifft_objectpsiK = pyfftw.FFTW(psiK, psiK,
flags=self.fft_object.flags,
axes=(0, 1), threads=self.N_THREADS,
direction="FFTW_BACKWARD")
psiK[:] = self.psi
assert pyfftw.is_n_byte_aligned(self.psi, self.al)
assert pyfftw.is_n_byte_aligned(psiK, self.al)
for step in xrange(N_TIMESTEPS):
# Record energy
if step % stepsPerObservation == 0 and recordEnergy:
density = np.absolute(self.psi) ** 2
gradx = np.gradient(self.psi)[0]
normFactor = density.sum()
self.energyTimes[step // stepsPerObservation] = self.time
self.energy[step // stepsPerObservation] = \
-(0.25 * np.gradient(
np.gradient(density)[0])[0]
- 0.5 * np.absolute(gradx) ** 2
- (self.g_C * density + self.g_R * n)
* density).sum() / normFactor
psiK[:] = self.psi
dK[:] = np.absolute(psiK) ** 2
assert pyfftw.is_n_byte_aligned(psiK, self.al)
diffusionUpdate(self.psi, self.fft_object, self.ifft_object,
kineticFactorRK4)
psiI[:] = self.psi
# K=1 step
# Updates on psiK and nK
nonLinearUpdatePsi(psiK, potential, dK,
n, dt)
diffusionUpdate(psiK, fft_objectpsiK, ifft_objectpsiK,
kineticFactorRK4)
# Update the next value of psi and n
self.psi += psiK / 6
# Update psik and nK. Store things for the next step
psiK /= 2
psiK += psiI
dK[:] = np.absolute(psiK) ** 2
self.time += dt / 2
# K=2 step
# Do updates
nonLinearUpdatePsi(psiK, potential, dK, n, dt)
# Update the next value of psi and n
self.psi += psiK / 3
# Update psiK and nK. Store things for the next step
psiK /= 2
psiK += psiI
dK[:] = np.absolute(psiK) ** 2
# K=3 step
# Do the updates
nonLinearUpdatePsi(psiK, potential, dK,
n, dt)
# Update the next value of psi and n
self.psi += psiK / 3
# Update psiK and nK. Store things for the next step
psiK += psiI
dK[:] = np.absolute(psiK) ** 2
self.time += dt / 2
# K=4 step
# Do the updates
diffusionUpdate(psiK, fft_objectpsiK, ifft_objectpsiK,
kineticFactorRK4)
nonLinearUpdatePsi(psiK, potential, dK, n, dt)
# Transform psi back
diffusionUpdate(self.psi, self.fft_object, self.ifft_object,
kineticFactorRK4)
# Update the next value of psi and n
self.psi += psiK / 6
dK[:] = np.absolute(self.psi) ** 2
# Do "exact" update on n
nonLinearUpdateN(n, P, dK, dt)
# Apply damping
self.psi *= self.damping
n *= self.damping
print("Time: %f" % self.time)
else:
raise ValueError("Invalid method")
self.nFinal = np.copy(n)
self.psiFinal = np.copy(self.psi)
def __init__(self, shape, dtype_in=None, data_in=None,
overwrite=True, inverse=True, packed=True, use_pyfftw=True):
try:
nx, ny, nz = shape
except (TypeError, ValueError):
raise ValueError('Expected 3D shape.')
if nx % 2 or ny % 2 or nz % 2:
raise ValueError('All shape dimensions must be even.')
if data_in is not None:
if not isinstance(data_in, np.ndarray):
raise ValueError(
'Invalid type for data_in: {0}.'.format(type(data_in)))
dtype_in = data_in.dtype
# Convert dtype_in to an object in the numpy scalar type hierarchy.
dtype_in = np.obj2sctype(dtype_in)
if dtype_in is None:
raise ValueError('Invalid dtype_in: {0}.'.format(dtype_in))
# Determine the input and output array type and shape.
if packed:
if inverse:
shape_in = (nx, ny, nz//2 + 1)
if not issubclass(dtype_in, np.complexfloating):
raise ValueError(
'Invalid dtype_in for inverse packed transform ' +
'(should be complex): {0}.'.format(dtype_in))
dtype_out = scalar_type(dtype_in)
shape_out = (nx, ny, nz + 2) if overwrite else shape
else:
shape_in = (nx, ny, nz + 2) if overwrite else shape
if not issubclass(dtype_in, np.floating):
raise ValueError(
'Invalid dtype_in for forward packed transform ' +
'(should be floating): {0}.'.format(dtype_in))
dtype_out = complex_type(dtype_in)
shape_out = (nx, ny, nz//2 + 1)
else:
if not issubclass(dtype_in, np.complexfloating):
raise ValueError(
'Expected complex dtype_in for transform: {0}.'
.format(dtype_in))
shape_in = shape_out = shape
dtype_out = dtype_in
if data_in is not None:
if data_in.shape != shape_in:
raise ValueError(
'data_in has wrong shape {0}, expected {1}.'
.format(data_in.shape, shape_in))
self.data_in = data_in
self.nbytes_allocated = 0
else:
# Allocate the input and output data buffers.
self.data_in = allocate(
shape_in, dtype_in, use_pyfftw=use_pyfftw)
self.nbytes_allocated = self.data_in.nbytes
if overwrite:
if packed:
# See https://github.com/hgomersall/pyFFTW/issues/29
self.data_out = self.data_in.view(dtype_out).reshape(shape_out)
# Hide the padding without copying. See http://www.fftw.org/doc/
# Multi_002dDimensional-DFTs-of-Real-Data.html.
if inverse:
self.data_out_padded = self.data_out
self.data_out = self.data_out[:, :, :nz]
else:
self.data_in_padded = self.data_in
self.data_in = self.data_in[:, :, :nz]
else:
self.data_out = self.data_in
else:
self.data_out = allocate(
shape_out, dtype_out, use_pyfftw=use_pyfftw)
self.nbytes_allocated += self.data_out.nbytes
# Try to use pyFFTW to configure the transform, if requested.
self.use_pyfftw = use_pyfftw
if self.use_pyfftw:
try:
import pyfftw
if not pyfftw.is_n_byte_aligned(self.data_in,
pyfftw.simd_alignment):
raise ValueError('data_in is not SIMD aligned.')
if not pyfftw.is_n_byte_aligned(self.data_out,
pyfftw.simd_alignment):
raise ValueError('data_out is not SIMD aligned.')
direction = 'FFTW_BACKWARD' if inverse else 'FFTW_FORWARD'
self.fftw_plan = pyfftw.FFTW(
self.data_in, self.data_out, direction=direction,
flags=('FFTW_ESTIMATE',), axes=(0, 1, 2))
self.fftw_norm = np.float(nx * ny * nz if inverse else 1)
except ImportError:
self.use_pyfftw = False
# Fall back to numpy.fft if we are not using pyFFTW.
if not self.use_pyfftw:
if inverse:
self.transformer = np.fft.irfftn if packed else np.fft.ifftn
else:
self.transformer = np.fft.rfftn if packed else np.fft.fftn
# Remember our options so we can create a reverse plan.
self.shape = shape
self.inverse = inverse
self.packed = packed
self.overwrite = overwrite