sumpy/examples/expansion-toys.py at b600ba1a15ca007d2f0049bb105e248cb2cd0414 · inducer/sumpy · GitHub

1
2
3
4
5
6
7
8
9
10
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
import numpy as np

import pyopencl as cl

import sumpy.toys as t
from sumpy.visualization import FieldPlotter
from sumpy.kernel import (      # noqa: F401
        YukawaKernel,
        HelmholtzKernel,
        LaplaceKernel)

try:
    import matplotlib.pyplot as plt
    USE_MATPLOTLIB = True
except ImportError:
    USE_MATPLOTLIB = False


def main():
    from sumpy.array_context import PyOpenCLArrayContext
    ctx = cl.create_some_context()
    queue = cl.CommandQueue(ctx)
    actx = PyOpenCLArrayContext(queue, force_device_scalars=True)

    tctx = t.ToyContext(
            # LaplaceKernel(2),
            YukawaKernel(2), extra_kernel_kwargs={"lam": 5},
            # HelmholtzKernel(2), extra_kernel_kwargs={"k": 0.3},
            )

    pt_src = t.PointSources(
            tctx,
            np.random.rand(2, 50) - 0.5,
            np.ones(50))

    fp = FieldPlotter([3, 0], extent=8)

    if USE_MATPLOTLIB:
        t.logplot(actx, fp, pt_src, cmap="jet")
        plt.colorbar()
        plt.show()

    mexp = t.multipole_expand(actx, pt_src, [0, 0], order=5)
    mexp2 = t.multipole_expand(actx, mexp, [0, 0.25])  # noqa: F841
    lexp = t.local_expand(actx, mexp, [3, 0])
    lexp2 = t.local_expand(actx, lexp, [3, 1], order=3)

    # diff = mexp - pt_src
    # diff = mexp2 - pt_src
    diff = lexp2 - pt_src

    print(t.l_inf(actx, diff, 1.2, center=lexp2.center))
    if USE_MATPLOTLIB:
        t.logplot(actx, fp, diff, cmap="jet", vmin=-3, vmax=0)
        plt.colorbar()
        plt.show()


if __name__ == "__main__":
    main()