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