def get_context(): """Set up context for classical (NS) solver""" float, complex, mpitype = datatypes(params.precision) collapse_fourier = False if params.dealias == '3/2-rule' else True dim = len(params.N) dtype = lambda d: float if d == dim - 1 else complex V = [ FunctionSpace(params.N[i], 'F', domain=(0, params.L[i]), dtype=dtype(i)) for i in range(dim) ] kw0 = { 'threads': params.threads, 'planner_effort': params.planner_effort['fft'] } T = TensorProductSpace(comm, V, dtype=float, slab=(params.decomposition == 'slab'), collapse_fourier=collapse_fourier, **kw0) VT = VectorSpace(T) # Different bases for nonlinear term, either 2/3-rule or 3/2-rule kw = { 'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1, 'dealias_direct': params.dealias == '2/3-rule' } Vp = [ FunctionSpace(params.N[i], 'F', domain=(0, params.L[i]), dtype=dtype(i), **kw) for i in range(dim) ] Tp = TensorProductSpace(comm, Vp, dtype=float, slab=(params.decomposition == 'slab'), collapse_fourier=collapse_fourier, **kw0) VTp = VectorSpace(Tp) mask = T.get_mask_nyquist() if params.mask_nyquist else None # Mesh variables X = T.local_mesh(True) K = T.local_wavenumbers(scaled=True) for i in range(dim): X[i] = X[i].astype(float) K[i] = K[i].astype(float) K2 = np.zeros(T.shape(True), dtype=float) for i in range(dim): K2 += K[i] * K[i] K_over_K2 = np.zeros(VT.shape(True), dtype=float) for i in range(dim): K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2) # Velocity and pressure. Use ndarray view for efficiency U = Array(VT) U_hat = Function(VT) P = Array(T) P_hat = Function(T) u_dealias = Array(VTp) # Primary variable u = U_hat # RHS array dU = Function(VT) curl = Array(VT) Source = Function(VT) # Possible source term initialized to zero work = work_arrays() hdf5file = NSFile(config.params.solver, checkpoint={ 'space': VT, 'data': { '0': { 'U': [U_hat] } } }, results={ 'space': VT, 'data': { 'U': [U], 'P': [P] } }) return config.AttributeDict(locals())
def get_context(): float, complex, mpitype = datatypes(params.precision) collapse_fourier = False if params.dealias == '3/2-rule' else True dim = len(params.N) dtype = lambda d: float if d == dim - 1 else complex V = [ FunctionSpace(params.N[i], 'F', domain=(0, params.L[i]), dtype=dtype(i)) for i in range(dim) ] kw0 = { 'threads': params.threads, 'planner_effort': params.planner_effort['fft'] } T = TensorProductSpace(comm, V, dtype=float, slab=(params.decomposition == 'slab'), collapse_fourier=collapse_fourier, **kw0) VT = VectorSpace(T) VM = CompositeSpace([T] * 2 * dim) mask = T.get_mask_nyquist() if params.mask_nyquist else None kw = { 'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1, 'dealias_direct': params.dealias == '2/3-rule' } Vp = [ FunctionSpace(params.N[i], 'F', domain=(0, params.L[i]), dtype=dtype(i), **kw) for i in range(dim) ] Tp = TensorProductSpace(comm, Vp, dtype=float, slab=(params.decomposition == 'slab'), collapse_fourier=collapse_fourier, **kw0) VTp = VectorSpace(Tp) VMp = CompositeSpace([Tp] * 2 * dim) # Mesh variables X = T.local_mesh(True) K = T.local_wavenumbers(scaled=True) for i in range(dim): X[i] = X[i].astype(float) K[i] = K[i].astype(float) K2 = np.zeros(T.shape(True), dtype=float) for i in range(dim): K2 += K[i] * K[i] K_over_K2 = np.zeros(VT.shape(True), dtype=float) for i in range(dim): K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2) UB = Array(VM) P = Array(T) curl = Array(VT) UB_hat = Function(VM) P_hat = Function(T) dU = Function(VM) Source = Array(VM) ub_dealias = Array(VMp) ZZ_hat = np.zeros((3, 3) + Tp.shape(True), dtype=complex) # Work array # Create views into large data structures U = UB[:3] U_hat = UB_hat[:3] B = UB[3:] B_hat = UB_hat[3:] # Primary variable u = UB_hat hdf5file = MHDFile(config.params.solver, checkpoint={ 'space': VM, 'data': { '0': { 'UB': [UB_hat] } } }, results={ 'space': VM, 'data': { 'UB': [UB] } }) return config.AttributeDict(locals())