finitedifferences.py 15.8 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
2
from typing import Optional, Union

3
import numpy as np
4
import sympy as sp
Martin Bauer's avatar
Martin Bauer committed
5

6
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
10
from pystencils.field import Field
from pystencils.simp.assignment_collection import AssignmentCollection
Martin Bauer's avatar
Martin Bauer committed
11
from pystencils.sympyextensions import fast_subs
12

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

15

16
# --------------------------------------- Advection Diffusion ----------------------------------------------------------
17
18


19
def diffusion(scalar, diffusion_coeff, idx=None):
20
21
22
23
24
25
26
    """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
27
        (f_W*d + f_S*d - 4*f_C*d + f_N*d + f_E*d)/dx**2
28
    """
29
30
31
32
33
34
    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")
35

36
37
38
39
    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
40

41

42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
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)


60
def transient(scalar, idx=None):
61
    """Transient term ∂_t(scalar)"""
62
63
64
65
66
67
68
69
70
    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)
71

72

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

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

Martin Bauer's avatar
Martin Bauer committed
86
    def _discretize_diffusion(self, e):
87
        result = 0
Martin Bauer's avatar
Martin Bauer committed
88
89
90
91
        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))
92
93
94
                           for offset in [-1, 1]]
            result += first_diffs[1] - first_diffs[0]
        return result / (self.dx ** 2)
95

96
97
98
    def _discretize_advection(self, expr):
        result = 0
        for c in range(expr.dim):
Martin Bauer's avatar
Martin Bauer committed
99
100
            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
101
102
103
                            for offset in [-1, 1]]
            result += interpolated[1] - interpolated[0]
        return result / self.dx
104

105
106
107
108
109
110
    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):
111
            arg, *indices = diff_args(e)
112
113
114
            from pystencils.interpolation_astnodes import InterpolatorAccess

            if not isinstance(arg, (Field.Access, InterpolatorAccess)):
115
116
                raise ValueError("Only derivatives with field or field accesses as arguments can be discretized")
            return self.spatial_stencil(indices, self.dx, arg)
117
118
119
        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
120

121
    def _discretize_diff(self, e):
122
        order = self._diff_order(e)
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
        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
143

144
145
146
    def __call__(self, expr):
        if isinstance(expr, list):
            return [self(e) for e in expr]
147
        elif isinstance(expr, sp.Matrix) or isinstance(expr, sp.ImmutableDenseMatrix):
148
149
150
151
            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
152

153
154
155
156
157
158
159
160
161
162
163
164
165
166
        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")
167
168


169
# -------------------------------------- Helper Classes ----------------------------------------------------------------
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184

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
185
    def scalar_index(self):
186
187
188
189
        return None if len(self.args) <= 2 else int(self.args[2])

    @property
    def dim(self):
Martin Bauer's avatar
Martin Bauer committed
190
        return self.scalar.spatial_dimensions
191
192

    def _latex(self, printer):
Martin Bauer's avatar
Martin Bauer committed
193
        name_suffix = "_%s" % self.scalar_index if self.scalar_index is not None else ""
194
195
        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
196
                                             printer.doprint(sp.Symbol(self.scalar.name + name_suffix)))
197
        else:
Martin Bauer's avatar
Martin Bauer committed
198
            args = [r"\partial_%d(%s %s)" % (i, printer.doprint(sp.Symbol(self.scalar.name + name_suffix)),
199
200
201
202
                                             printer.doprint(self.vector[i]))
                    for i in range(self.dim)]
            return " + ".join(args)

203
    # --- Interface for discretization strategy
204

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

Martin Bauer's avatar
Martin Bauer committed
213
214
215
    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)
216
217


218
219
220
221
222
223
224
class Diffusion(sp.Function):

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

    @property
Martin Bauer's avatar
Martin Bauer committed
225
    def diffusion_coeff(self):
226
227
228
229
230
231
        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
232
    def scalar_index(self):
233
234
235
236
        return None if len(self.args) <= 2 else int(self.args[2])

    @property
    def dim(self):
Martin Bauer's avatar
Martin Bauer committed
237
        return self.scalar.spatial_dimensions
238
239

    def _latex(self, printer):
Martin Bauer's avatar
Martin Bauer committed
240
241
242
243
        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
244
                                       printer.doprint(sp.Symbol(self.scalar.name + name_suffix)))
245

246
247
    # --- Interface for discretization strategy

Martin Bauer's avatar
Martin Bauer committed
248
249
250
    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)
251

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

259
260
261
262

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

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

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


277
278
279
280
281
282
283
284
285
286
# -------------------------------------------- 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
287
288
289
    Args:
        var: symbol to take the gradient of
        dim: dimension (length) of the gradient vector
290
291
292
    """
    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
293
    else:
294
        return [sp.Symbol("%s^Delta^%d" % (var.name, i)) for i in range(dim)]
295
296


297
298
299
300
301
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
302
303
304
305
306
307

    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
308

309
310
311
312
313
314
315
316
    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
317
      f_C*(-f_W/2 + f_E/2)
318
319
320
321
322
323
324
325
326
327
    """
    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))})
328
    return fast_subs(term, substitutions)
Martin Bauer's avatar
Martin Bauer committed
329

330

331
332
333
334
335
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.
336

337
338
339
340
341
342
343
344
    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
345

346
347
348
349
350
351
352
353
354
355
356
357
358
    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
359

360
361
362
363
    substitutions = {}
    for symbols, field in symbols_to_field_dict.items():
        if not hasattr(symbols, "__getitem__"):
            symbols = [symbols]
Martin Bauer's avatar
Martin Bauer committed
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
397
398
399
        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
400
        (f_W + f_S + f_B - 6*f_C + f_T + f_N + f_E)/dx**2
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
    """
    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