derivation.py 14.6 KB
 Stephan Seitz committed Jan 16, 2020 1 ``````import itertools `````` Martin Bauer committed Dec 03, 2018 2 ``````from collections import defaultdict `````` Martin Bauer committed Jul 11, 2019 3 `````` `````` Markus Holzer committed Aug 06, 2019 4 ``````import numpy as np `````` Martin Bauer committed Aug 08, 2019 5 ``````import sympy as sp `````` Martin Bauer committed Jul 11, 2019 6 `````` `````` Martin Bauer committed Dec 03, 2018 7 ``````from pystencils.field import Field `````` Michael Kuron committed Dec 03, 2019 8 ``````from pystencils.stencil import direction_string_to_offset `````` Martin Bauer committed Jul 11, 2019 9 10 ``````from pystencils.sympyextensions import multidimensional_sum, prod from pystencils.utils import LinearEquationSystem, fully_contains `````` Martin Bauer committed Dec 03, 2018 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 `````` class FiniteDifferenceStencilDerivation: """Derives finite difference stencils. Can derive standard finite difference stencils, as well as isotropic versions (see Isotropic Finite Differences by A. Kumar) Args: derivative_coordinates: tuple indicating which derivative should be approximated, (1, ) stands for first derivative in second direction (y), (0, 1) would be a mixed second derivative in x and y (0, 0, 0) would be a third derivative in x direction stencil: list of offset tuples, defining the stencil dx: spacing between grid points, one for all directions, i.e. dx=dy=dz Examples: Central differences >>> fd_1d = FiniteDifferenceStencilDerivation((0,), stencil=[(-1,), (0,), (1,)]) >>> result = fd_1d.get_stencil() >>> result Finite difference stencil of accuracy 2, isotropic error: False >>> result.weights [-1/2, 0, 1/2] Forward differences >>> fd_1d = FiniteDifferenceStencilDerivation((0,), stencil=[(0,), (1,)]) >>> result = fd_1d.get_stencil() >>> result Finite difference stencil of accuracy 1, isotropic error: False >>> result.weights [-1, 1] """ def __init__(self, derivative_coordinates, stencil, dx=1): self.dim = len(stencil[0]) self.field = Field.create_generic('f', spatial_dimensions=self.dim) self._derivative = tuple(sorted(derivative_coordinates)) self._stencil = stencil self._dx = dx self.weights = {tuple(d): self.symbolic_weight(*d) for d in self._stencil} def assume_symmetric(self, dim, anti_symmetric=False): """Adds restriction that weight in opposite directions of a dimension are equal (symmetric) or the negative of each other (anti symmetric) For example: dim=1, assumes that w(1, 1) == w(1, -1), if anti_symmetric=False or w(1, 1) == -w(1, -1) if anti_symmetric=True """ update = {} for direction, value in self.weights.items(): inv_direction = tuple(-offset if i == dim else offset for i, offset in enumerate(direction)) if direction[dim] < 0: inv_weight = self.weights[inv_direction] update[direction] = -inv_weight if anti_symmetric else inv_weight self.weights.update(update) def set_weight(self, offset, value): assert offset in self.weights self.weights[offset] = value def get_stencil(self, isotropic=False) -> 'FiniteDifferenceStencilDerivation.Result': weights = [self.weights[d] for d in self._stencil] system = LinearEquationSystem(sp.Matrix(weights).atoms(sp.Symbol)) order = 0 while True: new_system = system.copy() eq = self.error_term_equations(order) new_system.add_equations(eq) sol_structure = new_system.solution_structure() if sol_structure == 'single': system = new_system elif sol_structure == 'multiple': system = new_system elif sol_structure == 'none': break else: assert False order += 1 accuracy = order - len(self._derivative) error_is_isotropic = False if isotropic: new_system = system.copy() new_system.add_equations(self.isotropy_equations(order)) sol_structure = new_system.solution_structure() error_is_isotropic = sol_structure != 'none' if error_is_isotropic: system = new_system solve_res = system.solution() weight_list = [self.weights[d].subs(solve_res) for d in self._stencil] return self.Result(self._stencil, weight_list, accuracy, error_is_isotropic) @staticmethod def symbolic_weight(*args): str_args = [str(e) for e in args] `````` Markus Holzer committed Jun 19, 2020 110 `````` return sp.Symbol(f"w_({','.join(str_args)})") `````` Martin Bauer committed Dec 03, 2018 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 `````` def error_term_dict(self, order): error_terms = defaultdict(lambda: 0) for direction in self._stencil: weight = self.weights[tuple(direction)] x = tuple(self._dx * d_i for d_i in direction) for offset in multidimensional_sum(order, dim=self.field.spatial_dimensions): fac = sp.factorial(order) error_terms[tuple(sorted(offset))] += weight / fac * prod(x[off] for off in offset) if self._derivative in error_terms: error_terms[self._derivative] -= 1 return error_terms def error_term_equations(self, order): return list(self.error_term_dict(order).values()) def isotropy_equations(self, order): def cycle_int_sequence(sequence, modulus): result = [] arr = np.array(sequence, dtype=int) while True: if tuple(arr) in result: break result.append(tuple(arr)) arr = (arr + 1) % modulus return tuple(set(tuple(sorted(t)) for t in result)) error_dict = self.error_term_dict(order) eqs = [] for derivative_tuple in list(error_dict.keys()): if fully_contains(self._derivative, derivative_tuple): remaining = list(derivative_tuple) for e in self._derivative: del remaining[remaining.index(e)] permutations = cycle_int_sequence(remaining, self.dim) if len(permutations) == 1: eqs.append(error_dict[derivative_tuple]) else: for i in range(1, len(permutations)): `````` Martin Bauer committed Feb 18, 2019 150 151 `````` new_eq = (error_dict[tuple(sorted(permutations[i] + self._derivative))] - error_dict[tuple(sorted(permutations[i - 1] + self._derivative))]) `````` Martin Bauer committed Dec 03, 2018 152 153 154 155 156 157 158 159 160 161 162 163 164 165 `````` if new_eq: eqs.append(new_eq) else: eqs.append(error_dict[derivative_tuple]) return eqs class Result: def __init__(self, stencil, weights, accuracy, is_isotropic): self.stencil = stencil self.weights = weights self.accuracy = accuracy self.is_isotropic = is_isotropic def visualize(self): `````` Martin Bauer committed May 05, 2019 166 167 `````` from pystencils.stencil import plot plot(self.stencil, data=self.weights) `````` Martin Bauer committed Dec 03, 2018 168 169 170 171 172 `````` def apply(self, field_access: Field.Access): f = field_access return sum(f.get_shifted(*offset) * weight for offset, weight in zip(self.stencil, self.weights)) `````` Stephan Seitz committed Jan 16, 2020 173 174 175 `````` def __array__(self): return np.array(self.as_array().tolist()) `````` Markus Holzer committed Aug 06, 2019 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 `````` def as_array(self): dim = len(self.stencil[0]) assert (dim == 2 or dim == 3), "Only 2D or 3D matrix representations are available" max_offset = max(max(abs(e) for e in direction) for direction in self.stencil) shape_list = [] for i in range(dim): shape_list.append(2 * max_offset + 1) number_of_elements = np.prod(shape_list) shape = tuple(shape_list) result = sp.MutableDenseNDimArray([0] * number_of_elements, shape) if dim == 2: for direction, weight in zip(self.stencil, self.weights): result[max_offset - direction[1], max_offset + direction[0]] = weight if dim == 3: for direction, weight in zip(self.stencil, self.weights): result[max_offset - direction[1], max_offset + direction[0], max_offset + direction[2]] = weight return result `````` Stephan Seitz committed Jan 16, 2020 197 `````` def rotate_weights_and_apply(self, field_access: Field.Access, axes): `````` Markus Holzer committed Aug 06, 2019 198 199 200 201 `````` """derive gradient weights of other direction with already calculated weights of one direction via rotation and apply them to a field.""" dim = len(self.stencil[0]) assert (dim == 2 or dim == 3), "This function is only for 2D or 3D stencils available" `````` Markus Holzer committed Oct 13, 2020 202 `````` rotated_weights = np.rot90(np.array(self.__array__()), 1, axes) `````` Markus Holzer committed Aug 06, 2019 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 `````` result = [] max_offset = max(max(abs(e) for e in direction) for direction in self.stencil) if dim == 2: for direction in self.stencil: result.append(rotated_weights[max_offset - direction[1], max_offset + direction[0]]) if dim == 3: for direction in self.stencil: result.append(rotated_weights[max_offset - direction[1], max_offset + direction[0], max_offset + direction[2]]) f = field_access return sum(f.get_shifted(*offset) * weight for offset, weight in zip(self.stencil, result)) `````` Martin Bauer committed Dec 03, 2018 219 220 221 `````` def __repr__(self): return "Finite difference stencil of accuracy {}, isotropic error: {}".format(self.accuracy, self.is_isotropic) `````` Michael Kuron committed Dec 03, 2019 222 223 224 225 226 227 228 229 230 `````` class FiniteDifferenceStaggeredStencilDerivation: """Derives a finite difference stencil for application at a staggered position Args: neighbor: the neighbor direction string or vector at whose staggered position to calculate the derivative dim: how many dimensions (2 or 3) derivative: a tuple of directions over which to perform derivatives `````` Michael Kuron committed Jun 15, 2021 231 `````` free_weights_prefix: a string to prefix to free weight symbols. If None, do not return free weights `````` Michael Kuron committed Dec 03, 2019 232 233 `````` """ `````` Michael Kuron committed Jun 15, 2021 234 `````` def __init__(self, neighbor, dim, derivative=tuple(), free_weights_prefix=None): `````` Michael Kuron committed Dec 03, 2019 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 `````` 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: `````` Markus Holzer committed Oct 26, 2021 264 `````` main_points = [sp.Matrix(np.multiply(neighbor, sp.Matrix(c) / 2)) `````` Michael Kuron committed Dec 03, 2019 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 `````` 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])) `````` Michael Kuron committed Jun 15, 2021 285 286 287 288 `````` 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: `````` Michael Kuron committed Dec 03, 2019 289 290 291 292 293 294 295 296 297 298 `````` 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 `````` Michael Kuron committed Dec 06, 2019 299 `````` center = [tuple(p + pos) for p in points].index((0, 0, 0)[:dim]) `````` Michael Kuron committed Dec 03, 2019 300 `````` best = [b for b in best if b[center] != 0] `````` Michael Kuron committed Dec 06, 2019 301 `````` if len(best) > 1: # if there are still multiple, they are equivalent, so we average `````` Michael Kuron committed Dec 06, 2019 302 `````` weights = [sum([b[i] for b in best]) / len(best) for i in range(len(weights))] `````` Michael Kuron committed Dec 06, 2019 303 304 `````` else: weights = best[0] `````` Michael Kuron committed Dec 03, 2019 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 `````` assert weights points_tuple = tuple([tuple(p + pos) for p in points]) self._points = points_tuple self._weights = weights @property def points(self): """return the points of the stencil""" return self._points @property def stencil(self): """return the points of the stencil relative to the staggered position specified by neighbor""" return self._stencil @property def weights(self): """return the weights of the stencil""" assert self._weights is not None return self._weights def visualize(self): if self._weights is None: ws = None else: ws = np.array([w for w in self.weights if w != 0], dtype=float) pts = np.array([p for i, p in enumerate(self.points) if self.weights[i] != 0], dtype=int) from pystencils.stencil import plot plot(pts, data=ws) `````` Michael Kuron committed Dec 17, 2019 336 337 `````` def apply(self, access: Field.Access): return sum([access.get_shifted(*point) * weight for point, weight in zip(self.points, self.weights)])``````