Skip to content
Snippets Groups Projects
Select Git revision
  • master default
  • waldwiesen-simulation
  • weilin
  • textures
  • colab
  • torch_native-rewrite
  • ci-experiments
  • test-travis-test
  • os-travis
  • remove-lbm
  • bootstrap_theme
  • 0.3.3
  • 0.3.2
  • 0.3.1
  • 0.3.0
  • 0.2.2
  • 0.2.1
  • 0.2.0
  • 0.1.4
  • 0.1.3
  • 0.1.2
  • 0.1.1
  • 0.1.0
23 results
Blame Permalink
test_dynamic_function.py 3.08 KiB
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
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 






import sympy as sp

import pystencils
from pystencils.data_types import TypedSymbol, create_type
from pystencils_autodiff.framework_integration.astnodes import DynamicFunction
from pystencils_autodiff.framework_integration.printer import (
    DebugFrameworkPrinter, FrameworkIntegrationPrinter)
from pystencils_autodiff.framework_integration.types import TemplateType


def test_dynamic_function():
    x, y = pystencils.fields('x, y:  float32[3d]')

    a = sp.symbols('a')

    my_fun_call = DynamicFunction(TypedSymbol('my_fun',
                                              'std::function<double(double)>'), create_type('double'), a)

    assignments = pystencils.AssignmentCollection({
        y.center: x.center + my_fun_call
    })

    ast = pystencils.create_kernel(assignments)
    pystencils.show_code(ast, custom_backend=FrameworkIntegrationPrinter())

    template_fun_call = DynamicFunction(TypedSymbol('my_fun',
                                                    TemplateType('Functor_T')), create_type('double'), a, x.center)

    assignments = pystencils.AssignmentCollection({
        y.center: x.center + template_fun_call
    })

    ast = pystencils.create_kernel(assignments)
    pystencils.show_code(ast, custom_backend=FrameworkIntegrationPrinter())
    pystencils.show_code(ast, custom_backend=DebugFrameworkPrinter())


def test_dynamic_matrix():
    x, y = pystencils.fields('x, y:  float32[3d]')
    from pystencils.data_types import TypedMatrixSymbol

    a = sp.symbols('a')

    A = TypedMatrixSymbol('A', 3, 1, create_type('double'), 'Vector3<double>')

    my_fun_call = DynamicFunction(TypedSymbol('my_fun',
                                              'std::function<Vector3<double>(double)>'), A.dtype, a)

    assignments = pystencils.AssignmentCollection({
        A:  my_fun_call,
        y.center: A[0] + A[1] + A[2]
    })

    ast = pystencils.create_kernel(assignments)
    pystencils.show_code(ast, custom_backend=FrameworkIntegrationPrinter())


def test_dynamic_matrix_location_dependent():
    x, y = pystencils.fields('x, y:  float32[3d]')
    from pystencils.data_types import TypedMatrixSymbol

    A = TypedMatrixSymbol('A', 3, 1, create_type('double'), 'Vector3<double>')

    my_fun_call = DynamicFunction(TypedSymbol('my_fun',
                                              'std: : function < Vector3 < double > (int, int, int) >'),
                                  A.dtype,
                                  *pystencils.x_vector(3))

    assignments = pystencils.AssignmentCollection({
        A:  my_fun_call,
        y.center: A[0] + A[1] + A[2]
    })

    ast = pystencils.create_kernel(assignments)
    pystencils.show_code(ast, custom_backend=FrameworkIntegrationPrinter())

    my_fun_call = DynamicFunction(TypedSymbol('my_fun',
                                              TemplateType('Functor_T')), A.dtype, *pystencils.x_vector(3))

    assignments = pystencils.AssignmentCollection({
        A:  my_fun_call,
        y.center: A[0] + A[1] + A[2]
    })

    ast = pystencils.create_kernel(assignments)
    pystencils.show_code(ast, custom_backend=FrameworkIntegrationPrinter())