finitedifferences.py 15.6 KB
Newer Older
1
import numpy as np
2
import sympy as sp
3
from typing import Union, Optional
Martin Bauer's avatar
Martin Bauer committed
4

5
6
from pystencils import Field, AssignmentCollection
from pystencils.fd import Diff
7
8
from pystencils.fd.derivative import diff_args
from pystencils.fd.spatial import fd_stencils_standard
Martin Bauer's avatar
Martin Bauer committed
9
from pystencils.sympyextensions import fast_subs
10
11


12
FieldOrFieldAccess = Union[Field, Field.Access]
Martin Bauer's avatar
Martin Bauer committed
13

14

15
# --------------------------------------- Advection Diffusion ----------------------------------------------------------
16
17


18
def diffusion(scalar, diffusion_coeff, idx=None):
19
20
21
22
23
24
25
    """Diffusion term ∇·( diffusion_coeff · ∇(scalar))

    Examples:
        >>> f = Field.create_generic('f', spatial_dimensions=2)
        >>> diffusion_term = diffusion(scalar=f, diffusion_coeff=sp.Symbol("d"))
        >>> discretization = Discretization2ndOrder()
        >>> discretization(diffusion_term)
Martin Bauer's avatar
Martin Bauer committed
26
        (f_W*d + f_S*d - 4*f_C*d + f_N*d + f_E*d)/dx**2
27
    """
28
29
30
31
32
33
    if isinstance(scalar, Field):
        first_arg = scalar.center
    elif isinstance(scalar, Field.Access):
        first_arg = scalar
    else:
        raise ValueError("Diffused scalar has to be a pystencils Field or Field.Access")
34

35
36
37
38
    args = [first_arg, diffusion_coeff if not isinstance(diffusion_coeff, Field) else diffusion_coeff.center]
    if idx is not None:
        args.append(idx)
    return Diffusion(*args)
Martin Bauer's avatar
Martin Bauer committed
39

40

41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def advection(advected_scalar: FieldOrFieldAccess, velocity_field: FieldOrFieldAccess, idx: Optional[int] = None):
    """Advection term  ∇·(velocity_field · advected_scalar)

    Term that describes the advection of a scalar quantity in a velocity field.
    """
    if isinstance(advected_scalar, Field):
        first_arg = advected_scalar.center
    elif isinstance(advected_scalar, Field.Access):
        first_arg = advected_scalar
    else:
        raise ValueError("Advected scalar has to be a pystencils Field or Field.Access")

    args = [first_arg, velocity_field if not isinstance(velocity_field, Field) else velocity_field.center]
    if idx is not None:
        args.append(idx)
    return Advection(*args)


59
def transient(scalar, idx=None):
60
    """Transient term ∂_t(scalar)"""
61
62
63
64
65
66
67
68
69
    if isinstance(scalar, Field):
        args = [scalar.center]
    elif isinstance(scalar, Field.Access):
        args = [scalar]
    else:
        raise ValueError("Scalar has to be a pystencils Field or Field.Access")
    if idx is not None:
        args.append(idx)
    return Transient(*args)
70

71

72
class Discretization2ndOrder:
73
    def __init__(self, dx=sp.Symbol("dx"), dt=sp.Symbol("dt"), discretization_stencil_func=fd_stencils_standard):
74
75
        self.dx = dx
        self.dt = dt
76
        self.spatial_stencil = discretization_stencil_func
77

78
    @staticmethod
79
    def _diff_order(e):
80
81
82
        if not isinstance(e, Diff):
            return 0
        else:
83
            return 1 + Discretization2ndOrder._diff_order(e.args[0])
84

Martin Bauer's avatar
Martin Bauer committed
85
    def _discretize_diffusion(self, e):
86
        result = 0
Martin Bauer's avatar
Martin Bauer committed
87
88
89
90
        for c in range(e.dim):
            first_diffs = [offset
                           * (e.diffusion_scalar_at_offset(c, offset) * e.diffusion_coefficient_at_offset(c, offset)
                              - e.diffusion_scalar_at_offset(0, 0) * e.diffusion_coefficient_at_offset(0, 0))
91
92
93
                           for offset in [-1, 1]]
            result += first_diffs[1] - first_diffs[0]
        return result / (self.dx ** 2)
94

95
96
97
    def _discretize_advection(self, expr):
        result = 0
        for c in range(expr.dim):
Martin Bauer's avatar
Martin Bauer committed
98
99
            interpolated = [(expr.advected_scalar_at_offset(c, offset) * expr.velocity_field_at_offset(c, offset, c)
                             + expr.advected_scalar_at_offset(c, 0) * expr.velocity_field_at_offset(c, 0, c)) / 2
100
101
102
                            for offset in [-1, 1]]
            result += interpolated[1] - interpolated[0]
        return result / self.dx
103

104
105
106
107
108
109
    def _discretize_spatial(self, e):
        if isinstance(e, Diffusion):
            return self._discretize_diffusion(e)
        elif isinstance(e, Advection):
            return self._discretize_advection(e)
        elif isinstance(e, Diff):
110
111
112
113
            arg, *indices = diff_args(e)
            if not isinstance(arg, Field.Access):
                raise ValueError("Only derivatives with field or field accesses as arguments can be discretized")
            return self.spatial_stencil(indices, self.dx, arg)
114
115
116
        else:
            new_args = [self._discretize_spatial(a) for a in e.args]
            return e.func(*new_args) if new_args else e
Martin Bauer's avatar
Martin Bauer committed
117

118
    def _discretize_diff(self, e):
119
        order = self._diff_order(e)
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
        if order == 1:
            fa = e.args[0]
            index = e.target
            return (fa.neighbor(index, 1) - fa.neighbor(index, -1)) / (2 * self.dx)
        elif order == 2:
            indices = sorted([e.target, e.args[0].target])
            fa = e.args[0].args[0]
            if indices[0] == indices[1] and all(i >= 0 for i in indices):
                result = (-2 * fa + fa.neighbor(indices[0], -1) + fa.neighbor(indices[0], +1))
            elif indices[0] == indices[1]:
                result = 0
                for d in range(fa.field.spatial_dimensions):
                    result += (-2 * fa + fa.neighbor(d, -1) + fa.neighbor(d, +1))
            else:
                assert all(i >= 0 for i in indices)
                offsets = [(1, 1), [-1, 1], [1, -1], [-1, -1]]
                result = sum(o1 * o2 * fa.neighbor(indices[0], o1).neighbor(indices[1], o2) for o1, o2 in offsets) / 4
            return result / (self.dx ** 2)
        else:
            raise NotImplementedError("Term contains derivatives of order > 2")
Martin Bauer's avatar
Martin Bauer committed
140

141
142
143
    def __call__(self, expr):
        if isinstance(expr, list):
            return [self(e) for e in expr]
144
        elif isinstance(expr, sp.Matrix) or isinstance(expr, sp.ImmutableDenseMatrix):
145
146
147
148
            return expr.applyfunc(self.__call__)
        elif isinstance(expr, AssignmentCollection):
            return expr.copy(main_assignments=[e for e in expr.main_assignments],
                             subexpressions=[e for e in expr.subexpressions])
Martin Bauer's avatar
Martin Bauer committed
149

150
151
152
153
154
155
156
157
158
159
160
161
162
163
        transient_terms = expr.atoms(Transient)
        if len(transient_terms) == 0:
            return self._discretize_spatial(expr)
        elif len(transient_terms) == 1:
            transient_term = transient_terms.pop()
            solve_result = sp.solve(expr, transient_term)
            if len(solve_result) != 1:
                raise ValueError("Could not solve for transient term" + str(solve_result))
            rhs = solve_result.pop()
            # explicit euler
            return transient_term.scalar + self.dt * self._discretize_spatial(rhs)
        else:
            print(transient_terms)
            raise NotImplementedError("Cannot discretize expression with more than one transient term")
164
165


166
# -------------------------------------- Helper Classes ----------------------------------------------------------------
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181

class Advection(sp.Function):

    @property
    def scalar(self):
        return self.args[0].field

    @property
    def vector(self):
        if isinstance(self.args[1], Field.Access):
            return self.args[1].field
        else:
            return self.args[1]

    @property
Martin Bauer's avatar
Martin Bauer committed
182
    def scalar_index(self):
183
184
185
186
        return None if len(self.args) <= 2 else int(self.args[2])

    @property
    def dim(self):
Martin Bauer's avatar
Martin Bauer committed
187
        return self.scalar.spatial_dimensions
188
189

    def _latex(self, printer):
Martin Bauer's avatar
Martin Bauer committed
190
        name_suffix = "_%s" % self.scalar_index if self.scalar_index is not None else ""
191
192
        if isinstance(self.vector, Field):
            return r"\nabla \cdot(%s %s)" % (printer.doprint(sp.Symbol(self.vector.name)),
Martin Bauer's avatar
Martin Bauer committed
193
                                             printer.doprint(sp.Symbol(self.scalar.name + name_suffix)))
194
        else:
Martin Bauer's avatar
Martin Bauer committed
195
            args = [r"\partial_%d(%s %s)" % (i, printer.doprint(sp.Symbol(self.scalar.name + name_suffix)),
196
197
198
199
                                             printer.doprint(self.vector[i]))
                    for i in range(self.dim)]
            return " + ".join(args)

200
    # --- Interface for discretization strategy
201

Martin Bauer's avatar
Martin Bauer committed
202
    def velocity_field_at_offset(self, offset_dim, offset_value, index):
203
204
        v = self.vector
        if isinstance(v, Field):
Martin Bauer's avatar
Martin Bauer committed
205
206
            assert v.index_dimensions == 1
            return v.neighbor(offset_dim, offset_value)(index)
207
208
209
        else:
            return v[index]

Martin Bauer's avatar
Martin Bauer committed
210
211
212
    def advected_scalar_at_offset(self, offset_dim, offset_value):
        idx = 0 if self.scalar_index is None else int(self.scalar_index)
        return self.scalar.neighbor(offset_dim, offset_value)(idx)
213
214


215
216
217
218
219
220
221
class Diffusion(sp.Function):

    @property
    def scalar(self):
        return self.args[0].field

    @property
Martin Bauer's avatar
Martin Bauer committed
222
    def diffusion_coeff(self):
223
224
225
226
227
228
        if isinstance(self.args[1], Field.Access):
            return self.args[1].field
        else:
            return self.args[1]

    @property
Martin Bauer's avatar
Martin Bauer committed
229
    def scalar_index(self):
230
231
232
233
        return None if len(self.args) <= 2 else int(self.args[2])

    @property
    def dim(self):
Martin Bauer's avatar
Martin Bauer committed
234
        return self.scalar.spatial_dimensions
235
236

    def _latex(self, printer):
Martin Bauer's avatar
Martin Bauer committed
237
238
239
240
        name_suffix = "_%s" % self.scalar_index if self.scalar_index is not None else ""
        coeff = self.diffusion_coeff
        diff_coeff = sp.Symbol(coeff.name) if isinstance(coeff, Field) else coeff
        return r"div(%s \nabla %s)" % (printer.doprint(diff_coeff),
Martin Bauer's avatar
Martin Bauer committed
241
                                       printer.doprint(sp.Symbol(self.scalar.name + name_suffix)))
242

243
244
    # --- Interface for discretization strategy

Martin Bauer's avatar
Martin Bauer committed
245
246
247
    def diffusion_scalar_at_offset(self, offset_dim, offset_value):
        idx = 0 if self.scalar_index is None else self.scalar_index
        return self.scalar.neighbor(offset_dim, offset_value)(idx)
248

Martin Bauer's avatar
Martin Bauer committed
249
250
    def diffusion_coefficient_at_offset(self, offset_dim, offset_value):
        d = self.diffusion_coeff
251
        if isinstance(d, Field):
Martin Bauer's avatar
Martin Bauer committed
252
            return d.neighbor(offset_dim, offset_value)
253
254
255
        else:
            return d

256
257
258
259

class Transient(sp.Function):
    @property
    def scalar(self):
Martin Bauer's avatar
Martin Bauer committed
260
        if self.scalar_index is None:
Martin Bauer's avatar
Martin Bauer committed
261
            return self.args[0].field.center
262
        else:
Martin Bauer's avatar
Martin Bauer committed
263
            return self.args[0].field(self.scalar_index)
264
265

    @property
Martin Bauer's avatar
Martin Bauer committed
266
    def scalar_index(self):
267
268
269
        return None if len(self.args) <= 1 else int(self.args[1])

    def _latex(self, printer):
Martin Bauer's avatar
Martin Bauer committed
270
        name_suffix = "_%s" % self.scalar_index if self.scalar_index is not None else ""
Martin Bauer's avatar
Martin Bauer committed
271
        return r"\partial_t %s" % (printer.doprint(sp.Symbol(self.scalar.name + name_suffix)),)
272
273


274
275
276
277
278
279
280
281
282
283
# -------------------------------------------- Deprecated Functions ----------------------------------------------------


def grad(var, dim=3):
    r"""
    Gradients are represented as a special symbol:
    e.g. :math:`\nabla x = (x^{\Delta 0}, x^{\Delta 1}, x^{\Delta 2})`

    This function takes a symbol and creates the gradient symbols according to convention above

Martin Bauer's avatar
Martin Bauer committed
284
285
286
    Args:
        var: symbol to take the gradient of
        dim: dimension (length) of the gradient vector
287
288
289
    """
    if hasattr(var, "__getitem__"):
        return [[sp.Symbol("%s^Delta^%d" % (v.name, i)) for v in var] for i in range(dim)]
Martin Bauer's avatar
Martin Bauer committed
290
    else:
291
        return [sp.Symbol("%s^Delta^%d" % (var.name, i)) for i in range(dim)]
292
293


294
295
296
297
298
def discretize_center(term, symbols_to_field_dict, dx, dim=3):
    """
    Expects term that contains given symbols and gradient components of these symbols and replaces them
    by field accesses. Gradients are replaced by centralized approximations:
    ``(upper neighbor - lower neighbor ) / ( 2*dx)``
Martin Bauer's avatar
Martin Bauer committed
299
300
301
302
303
304

    Args:
        term: term where symbols and gradient(symbol) should be replaced
        symbols_to_field_dict: mapping of symbols to Field
        dx: width and height of one cell
        dim: dimension
305

306
307
308
309
310
311
312
313
    Example:
      >>> x = sp.Symbol("x")
      >>> grad_x = grad(x, dim=3)
      >>> term = x * grad_x[0]
      >>> term
      x*x^Delta^0
      >>> f = Field.create_generic('f', spatial_dimensions=3)
      >>> discretize_center(term, { x: f }, dx=1, dim=3)
Martin Bauer's avatar
Martin Bauer committed
314
      f_C*(-f_W/2 + f_E/2)
315
316
317
318
319
320
321
322
323
324
    """
    substitutions = {}
    for symbols, field in symbols_to_field_dict.items():
        if not hasattr(symbols, "__getitem__"):
            symbols = [symbols]
        g = grad(symbols, dim)
        substitutions.update({symbol: field(i) for i, symbol in enumerate(symbols)})
        for d in range(dim):
            up, down = __up_down_offsets(d, dim)
            substitutions.update({g[d][i]: (field[up](i) - field[down](i)) / dx / 2 for i in range(len(symbols))})
325
    return fast_subs(term, substitutions)
Martin Bauer's avatar
Martin Bauer committed
326

327

328
329
330
331
332
def discretize_staggered(term, symbols_to_field_dict, coordinate, coordinate_offset, dx, dim=3):
    """
    Expects term that contains given symbols and gradient components of these symbols and replaces them
    by field accesses. Gradients in coordinate direction  are replaced by staggered version at cell boundary.
    Symbols themselves and gradients in other directions are replaced by interpolated version at cell face.
333

334
335
336
337
338
339
340
341
    Args:
        term: input term where symbols and gradients are replaced
        symbols_to_field_dict: mapping of symbols to Field
        coordinate: id for coordinate (0 for x, 1 for y, ... ) defining cell boundary.
                    Only gradients in this direction are replaced e.g. if symbol^Delta^coordinate
        coordinate_offset: either +1 or -1 for upper or lower face in coordinate direction
        dx: width and height of one cell
        dim: dimension
342

343
344
345
346
347
348
349
350
351
352
353
354
355
    Examples:
      Discretizing at right/east face of cell i.e. coordinate=0, offset=1)
      >>> x, dx = sp.symbols("x dx")
      >>> grad_x = grad(x, dim=3)
      >>> term = x * grad_x[0]
      >>> term
      x*x^Delta^0
      >>> f = Field.create_generic('f', spatial_dimensions=3)
      >>> discretize_staggered(term, symbols_to_field_dict={ x: f}, dx=dx, coordinate=0, coordinate_offset=1, dim=3)
      (-f_C + f_E)*(f_C/2 + f_E/2)/dx
    """
    assert coordinate_offset == 1 or coordinate_offset == -1
    assert 0 <= coordinate < dim
Martin Bauer's avatar
Martin Bauer committed
356

357
358
359
360
    substitutions = {}
    for symbols, field in symbols_to_field_dict.items():
        if not hasattr(symbols, "__getitem__"):
            symbols = [symbols]
Martin Bauer's avatar
Martin Bauer committed
361

362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
        offset = [0] * dim
        offset[coordinate] = coordinate_offset
        offset = np.array(offset, dtype=np.int)

        gradient = grad(symbols)[coordinate]
        substitutions.update({s: (field[offset](i) + field(i)) / 2 for i, s in enumerate(symbols)})
        substitutions.update({g: (field[offset](i) - field(i)) / dx * coordinate_offset
                              for i, g in enumerate(gradient)})
        for d in range(dim):
            if d == coordinate:
                continue
            up, down = __up_down_offsets(d, dim)
            for i, s in enumerate(symbols):
                center_grad = (field[up](i) - field[down](i)) / (2 * dx)
                neighbor_grad = (field[up + offset](i) - field[down + offset](i)) / (2 * dx)
                substitutions[grad(s)[d]] = (center_grad + neighbor_grad) / 2

    return fast_subs(term, substitutions)


def discretize_divergence(vector_term, symbols_to_field_dict, dx):
    """
    Computes discrete divergence of symbolic vector

    Args:
        vector_term: sequence of terms, interpreted as vector
        symbols_to_field_dict: mapping of symbols to Field
        dx: length of a cell

    Examples:
        Laplace stencil
        >>> x, dx = sp.symbols("x dx")
        >>> grad_x = grad(x, dim=3)
        >>> f = Field.create_generic('f', spatial_dimensions=3)
        >>> sp.simplify(discretize_divergence(grad_x, {x : f}, dx))
Martin Bauer's avatar
Martin Bauer committed
397
        (f_W + f_S + f_B - 6*f_C + f_T + f_N + f_E)/dx**2
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
    """
    dim = len(vector_term)
    result = 0
    for d in range(dim):
        for offset in [-1, 1]:
            result += offset * discretize_staggered(vector_term[d], symbols_to_field_dict, d, offset, dx, dim)
    return result / dx


def __up_down_offsets(d, dim):
    coord = [0] * dim
    coord[d] = 1
    up = np.array(coord, dtype=np.int)
    coord[d] = -1
    down = np.array(coord, dtype=np.int)
    return up, down