def __init__(self, normal_direction, stencil, name=None): if isinstance(normal_direction, str): normal_direction = direction_string_to_offset(normal_direction, dim=len(stencil[0])) if name is None: name = f"Simple Outflow: {offset_to_direction_string(normal_direction)}" self.normal_direction = normal_direction super(SimpleExtrapolationOutflow, self).__init__(name)
def __getitem__(self, offset): if type(offset) is np.ndarray: offset = tuple(offset) if type(offset) is str: offset = tuple( direction_string_to_offset(offset, self.spatial_dimensions)) if type(offset) is not tuple: offset = (offset, ) if len(offset) != self.spatial_dimensions: raise ValueError("Wrong number of spatial indices: " "Got %d, expected %d" % (len(offset), self.spatial_dimensions)) return Field.Access(self, offset)
def test_field_slice(): """Tests (un)packing slices of a scalar field (from)to a buffer.""" fields = _generate_fields() for d in ['N', 'S', 'NW', 'SW', 'TNW', 'B']: for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields: # Extract slice from N direction of the field slice_dir = direction_string_to_offset(d, dim=len(src_arr.shape)) pack_slice = get_slice_before_ghost_layer(slice_dir) unpack_slice = get_ghost_region_slice(slice_dir) src_field = Field.create_from_numpy_array("src_field", src_arr[pack_slice]) dst_field = Field.create_from_numpy_array("dst_field", src_arr[unpack_slice]) buffer = Field.create_generic("buffer", spatial_dimensions=1, field_type=FieldType.BUFFER, dtype=src_arr.dtype) pack_eqs = [Assignment(buffer.center(), src_field.center())] pack_types = { 'src_field': gpu_src_arr.dtype, 'buffer': gpu_buffer_arr.dtype } pack_code = create_cuda_kernel(pack_eqs, type_info=pack_types) pack_kernel = make_python_function(pack_code) pack_kernel(buffer=gpu_buffer_arr, src_field=gpu_src_arr[pack_slice]) # Unpack into ghost layer of dst_field in N direction unpack_eqs = [Assignment(dst_field.center(), buffer.center())] unpack_types = { 'dst_field': gpu_dst_arr.dtype, 'buffer': gpu_buffer_arr.dtype } unpack_code = create_cuda_kernel(unpack_eqs, type_info=unpack_types) unpack_kernel = make_python_function(unpack_code) unpack_kernel(buffer=gpu_buffer_arr, dst_field=gpu_dst_arr[unpack_slice]) dst_arr = gpu_dst_arr.get() np.testing.assert_equal(src_arr[pack_slice], dst_arr[unpack_slice])
def condition(direction): """exclude those staggered points that correspond to fluxes between ghost cells""" exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim]) for elementary_direction in direction: exclusions.remove(inverse_direction_string(elementary_direction)) conditions = [] for e in exclusions: if e in common_exclusions: continue offset = direction_string_to_offset(e) for i, o in enumerate(offset): if o == 1: conditions.append(counters[i] < shape[i] - 1) elif o == -1: conditions.append(counters[i] > 0) return sp.And(*conditions)
def __init__(self, normal_direction, lb_method, dt=1, dx=1, name=None, streaming_pattern='pull', zeroth_timestep=Timestep.BOTH, initial_density=None, initial_velocity=None, data_type='double'): self.lb_method = lb_method self.stencil = lb_method.stencil self.dim = len(self.stencil[0]) if isinstance(normal_direction, str): normal_direction = direction_string_to_offset(normal_direction, dim=self.dim) if name is None: name = f"Outflow: {offset_to_direction_string(normal_direction)}" self.normal_direction = normal_direction self.streaming_pattern = streaming_pattern self.zeroth_timestep = zeroth_timestep self.dx = sp.Number(dx) self.dt = sp.Number(dt) self.c = sp.sqrt(sp.Rational(1, 3)) * (self.dx / self.dt) self.initial_density = initial_density self.initial_velocity = initial_velocity self.equilibrium_calculation = None self.data_type = data_type if initial_density and initial_velocity: equilibrium = lb_method.get_equilibrium(conserved_quantity_equations=AssignmentCollection([])) rho = lb_method.zeroth_order_equilibrium_moment_symbol u_vec = lb_method.first_order_equilibrium_moment_symbols eq_lambda = equilibrium.lambdify((rho,) + u_vec) post_pdf_symbols = lb_method.post_collision_pdf_symbols def calc_eq_pdfs(density, velocity, j): return eq_lambda(density, *velocity)[post_pdf_symbols[j]] self.equilibrium_calculation = calc_eq_pdfs super(ExtrapolationOutflow, self).__init__(name)
def __init__(self, neighbor, dim, derivative=tuple(), free_weights_prefix=None): if type(neighbor) is str: neighbor = direction_string_to_offset(neighbor) if dim == 2: assert neighbor[dim:] == 0 assert derivative is tuple() or max(derivative) < dim neighbor = sp.Matrix(neighbor[:dim]) pos = neighbor / 2 def unitvec(i): """return the `i`-th unit vector in three dimensions""" a = np.zeros(dim, dtype=int) a[i] = 1 return a def flipped(a, i): """return `a` with its `i`-th element's sign flipped""" a = a.copy() a[i] *= -1 return a # determine the points to use, coordinates are relative to position points = [] if np.linalg.norm(neighbor, 1) == 1: main_points = [neighbor / 2, neighbor / -2] elif np.linalg.norm(neighbor, 1) == 2: nonzero_indices = [ i for i, v in enumerate(neighbor) if v != 0 and i < dim ] main_points = [ neighbor / 2, neighbor / -2, flipped(neighbor / 2, nonzero_indices[0]), flipped(neighbor / -2, nonzero_indices[0]) ] else: main_points = [ sp.Matrix(np.multiply(neighbor, sp.Matrix(c) / 2)) for c in itertools.product([-1, 1], repeat=3) ] points += main_points zero_indices = [ i for i, v in enumerate(neighbor) if v == 0 and i < dim ] for i in zero_indices: points += [point + sp.Matrix(unitvec(i)) for point in main_points] points += [point - sp.Matrix(unitvec(i)) for point in main_points] points_tuple = tuple([tuple(p) for p in points]) self._stencil = points_tuple # determine the stencil weights if len(derivative) == 0: weights = None else: derivation = FiniteDifferenceStencilDerivation( derivative, points_tuple).get_stencil() if not derivation.accuracy: raise Exception( 'the requested derivative cannot be performed with the available neighbors' ) weights = derivation.weights # if the weights are underdefined, we can choose the free symbols to find the sparsest stencil free_weights = set( itertools.chain(*[w.free_symbols for w in weights])) if free_weights_prefix is not None: weights = [ w.subs({ fw: sp.Symbol(f"{free_weights_prefix}_{i}") for i, fw in enumerate(free_weights) }) for w in weights ] elif len(free_weights) > 0: zero_counts = defaultdict(list) for values in itertools.product( [-1, -sp.Rational(1, 2), 0, 1, sp.Rational(1, 2)], repeat=len(free_weights)): subs = { free_weight: value for free_weight, value in zip(free_weights, values) } weights = [w.subs(subs) for w in derivation.weights] if not all(a == 0 for a in weights): zero_count = sum([1 for w in weights if w == 0]) zero_counts[zero_count].append(weights) best = zero_counts[max(zero_counts.keys())] if len( best ) > 1: # if there are multiple, pick the one that contains a nonzero center weight center = [tuple(p + pos) for p in points].index( (0, 0, 0)[:dim]) best = [b for b in best if b[center] != 0] if len( best ) > 1: # if there are still multiple, they are equivalent, so we average weights = [ sum([b[i] for b in best]) / len(best) for i in range(len(weights)) ] else: weights = best[0] assert weights points_tuple = tuple([tuple(p + pos) for p in points]) self._points = points_tuple self._weights = weights
def create_staggered_kernel(assignments, target: Target = Target.CPU, gpu_exclusive_conditions=False, **kwargs): """Kernel that updates a staggered field. .. image:: /img/staggered_grid.svg For a staggered field, the first index coordinate defines the location of the staggered value. Further index coordinates can be used to store vectors/tensors at each point. Args: assignments: a sequence of assignments or an AssignmentCollection. Assignments to staggered field are processed specially, while subexpressions and assignments to regular fields are passed through to `create_kernel`. Multiple different staggered fields can be used, but they all need to use the same stencil (i.e. the same number of staggered points) and shape. target: 'CPU' or 'GPU' gpu_exclusive_conditions: disable the use of multiple conditionals inside the loop. The outer layers are then handled in an else branch. kwargs: passed directly to create_kernel, iteration_slice and ghost_layers parameters are not allowed Returns: AST, see `create_kernel` """ if 'ghost_layers' in kwargs: assert kwargs['ghost_layers'] is None del kwargs['ghost_layers'] if 'iteration_slice' in kwargs: assert kwargs['iteration_slice'] is None del kwargs['iteration_slice'] if 'omp_single_loop' in kwargs: assert kwargs['omp_single_loop'] is False del kwargs['omp_single_loop'] if isinstance(assignments, AssignmentCollection): subexpressions = assignments.subexpressions + [ a for a in assignments.main_assignments if not hasattr(a, 'lhs') or type(a.lhs) is not Field.Access or not FieldType.is_staggered(a.lhs.field) ] assignments = [ a for a in assignments.main_assignments if hasattr(a, 'lhs') and type(a.lhs) is Field.Access and FieldType.is_staggered(a.lhs.field) ] else: subexpressions = [ a for a in assignments if not hasattr(a, 'lhs') or type(a.lhs) is not Field.Access or not FieldType.is_staggered(a.lhs.field) ] assignments = [ a for a in assignments if hasattr(a, 'lhs') and type(a.lhs) is Field.Access and FieldType.is_staggered(a.lhs.field) ] if len(set([tuple(a.lhs.field.staggered_stencil) for a in assignments])) != 1: raise ValueError( "All assignments need to be made to staggered fields with the same stencil" ) if len(set([a.lhs.field.shape for a in assignments])) != 1: raise ValueError( "All assignments need to be made to staggered fields with the same shape" ) staggered_field = assignments[0].lhs.field stencil = staggered_field.staggered_stencil dim = staggered_field.spatial_dimensions shape = staggered_field.shape counters = [ LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim) ] final_assignments = [] # find out whether any of the ghost layers is not needed common_exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim]) for direction in stencil: exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim]) for elementary_direction in direction: exclusions.remove(inverse_direction_string(elementary_direction)) common_exclusions.intersection_update(exclusions) ghost_layers = [[0, 0] for d in range(dim)] for direction in common_exclusions: direction = direction_string_to_offset(direction) for d, s in enumerate(direction): if s == 1: ghost_layers[d][1] = 1 elif s == -1: ghost_layers[d][0] = 1 def condition(direction): """exclude those staggered points that correspond to fluxes between ghost cells""" exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim]) for elementary_direction in direction: exclusions.remove(inverse_direction_string(elementary_direction)) conditions = [] for e in exclusions: if e in common_exclusions: continue offset = direction_string_to_offset(e) for i, o in enumerate(offset): if o == 1: conditions.append(counters[i] < shape[i] - 1) elif o == -1: conditions.append(counters[i] > 0) return sp.And(*conditions) if gpu_exclusive_conditions: outer_assignment = None conditions = {direction: condition(direction) for direction in stencil} for num_conditions in range(len(stencil)): for combination in itertools.combinations(conditions.values(), num_conditions): for assignment in assignments: direction = stencil[assignment.lhs.index[0]] if conditions[direction] in combination: assignment = SympyAssignment(assignment.lhs, assignment.rhs) outer_assignment = Conditional(sp.And(*combination), Block([assignment]), outer_assignment) inner_assignment = [] for assignment in assignments: inner_assignment.append( SympyAssignment(assignment.lhs, assignment.rhs)) last_conditional = Conditional( sp.And(*[condition(d) for d in stencil]), Block(inner_assignment), outer_assignment) final_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \ [SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \ [last_conditional] if target == Target.CPU: from pystencils.cpu import create_kernel as create_kernel_cpu ast = create_kernel_cpu(final_assignments, ghost_layers=ghost_layers, omp_single_loop=False, **kwargs) else: ast = create_kernel(final_assignments, ghost_layers=ghost_layers, target=target, **kwargs) return ast for assignment in assignments: direction = stencil[assignment.lhs.index[0]] sp_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \ [SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \ [SympyAssignment(assignment.lhs, assignment.rhs)] last_conditional = Conditional(condition(direction), Block(sp_assignments)) final_assignments.append(last_conditional) remove_start_conditional = any([gl[0] == 0 for gl in ghost_layers]) prepend_optimizations = [ lambda ast: remove_conditionals_in_staggered_kernel( ast, remove_start_conditional), move_constants_before_loop ] if 'cpu_prepend_optimizations' in kwargs: prepend_optimizations += kwargs['cpu_prepend_optimizations'] del kwargs['cpu_prepend_optimizations'] ast = create_kernel(final_assignments, ghost_layers=ghost_layers, target=target, omp_single_loop=False, cpu_prepend_optimizations=prepend_optimizations, **kwargs) return ast