test_tfmad.py 8.95 KB
Newer Older
Stephan Seitz's avatar
Stephan Seitz committed
1
import argparse
Stephan Seitz's avatar
Stephan Seitz committed
2
import os
Stephan Seitz's avatar
Stephan Seitz committed
3
4

import numpy as np
Stephan Seitz's avatar
Stephan Seitz committed
5
import pytest
Stephan Seitz's avatar
Stephan Seitz committed
6
7
8
9
import sympy as sp

import pystencils as ps
import pystencils_autodiff
Stephan Seitz's avatar
Stephan Seitz committed
10
from test_utils.gradient_check_tensorflow import compute_gradient_error_without_border
Stephan Seitz's avatar
Stephan Seitz committed
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


def test_tfmad_stencil():

    f, out = ps.fields("f, out: double[2D]")

    cont = ps.fd.Diff(f, 0) - ps.fd.Diff(f, 1)
    discretize = ps.fd.Discretization2ndOrder(dx=1)
    discretization = discretize(cont)

    assignment = ps.Assignment(out.center(), discretization)
    assignment_collection = ps.AssignmentCollection([assignment], [])
    print('Forward')
    print(assignment_collection)

    print('Backward')
    backward = pystencils_autodiff.create_backward_assignments(
        assignment_collection, diff_mode='transposed-forward')
    print(backward)


def test_tfmad_two_stencils():

    a, b, out = ps.fields("a, b, out: double[2D]")

    cont = ps.fd.Diff(a, 0) - ps.fd.Diff(a, 1) - \
        ps.fd.Diff(b, 0) + ps.fd.Diff(b, 1)
    discretize = ps.fd.Discretization2ndOrder(dx=1)
    discretization = discretize(cont)

    assignment = ps.Assignment(out.center(), discretization)
    assignment_collection = ps.AssignmentCollection([assignment], [])
    print('Forward')
    print(assignment_collection)

    print('Backward')
Stephan Seitz's avatar
Stephan Seitz committed
47
48
    auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection,
                                               diff_mode='transposed-forward')
Stephan Seitz's avatar
Stephan Seitz committed
49
50
51
52
53
    backward = auto_diff.backward_assignments
    print(backward)
    print('Forward output fields (to check order)')
    print(auto_diff.forward_input_fields)

Stephan Seitz's avatar
Stephan Seitz committed
54
55
    print(auto_diff)

Stephan Seitz's avatar
Stephan Seitz committed
56

57
@pytest.mark.skipif("CI" in os.environ, reason="Temporary skip")
Stephan Seitz's avatar
Stephan Seitz committed
58
def test_tfmad_gradient_check():
Stephan Seitz's avatar
Stephan Seitz committed
59
60
    tf = pytest.importorskip('tensorflow')

Stephan Seitz's avatar
Stephan Seitz committed
61
    a, b, out = ps.fields("a, b, out: double[21,13]")
62
    print(a.shape)
Stephan Seitz's avatar
Stephan Seitz committed
63

Stephan Seitz's avatar
Stephan Seitz committed
64
65
    cont = ps.fd.Diff(a, 0) - ps.fd.Diff(a, 1) - ps.fd.Diff(b, 0) + ps.fd.Diff(
        b, 1)
Stephan Seitz's avatar
Stephan Seitz committed
66
67
68
69
70
71
72
73
74
    discretize = ps.fd.Discretization2ndOrder(dx=1)
    discretization = discretize(cont)

    assignment = ps.Assignment(out.center(), discretization)
    assignment_collection = ps.AssignmentCollection([assignment], [])
    print('Forward')
    print(assignment_collection)

    print('Backward')
Stephan Seitz's avatar
Stephan Seitz committed
75
76
    auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection,
                                               diff_mode='transposed-forward')
Stephan Seitz's avatar
Stephan Seitz committed
77
78
79
80
81
82
83
84
85
86
87
88
89
    backward = auto_diff.backward_assignments
    print(backward)
    print('Forward output fields (to check order)')
    print(auto_diff.forward_input_fields)

    a_tensor = tf.Variable(np.zeros(a.shape, a.dtype.numpy_dtype))
    b_tensor = tf.Variable(np.zeros(a.shape, a.dtype.numpy_dtype))
    out_tensor = auto_diff.create_tensorflow_op({a: a_tensor, b: b_tensor})

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())

        gradient_error = compute_gradient_error_without_border(
Stephan Seitz's avatar
Stephan Seitz committed
90
91
92
93
94
            [a_tensor, b_tensor], [a.shape, b.shape],
            out_tensor,
            out.shape,
            num_border_pixels=2,
            ndim=2)
Stephan Seitz's avatar
Stephan Seitz committed
95
96
        print('error: %s' % gradient_error.max_error)

Stephan Seitz's avatar
Stephan Seitz committed
97
        assert any(e < 1e-4 for e in gradient_error.values())
Stephan Seitz's avatar
Stephan Seitz committed
98
99
100
101
102
103
104
105
106


def check_tfmad_vector_input_data(args):
    dtype = args.dtype
    domain_shape = args.domain_shape
    ndim = len(domain_shape)

    # create arrays
    c_arr = np.zeros(domain_shape)
Stephan Seitz's avatar
Stephan Seitz committed
107
    v_arr = np.zeros(domain_shape + (ndim, ))
Stephan Seitz's avatar
Stephan Seitz committed
108
109

    # create fields
Stephan Seitz's avatar
Stephan Seitz committed
110
111
112
    c, v, c_next = ps.fields("c, v(2), c_next: % s[ % i, % i]" %
                             ("float" if dtype == np.float32 else "double",
                              domain_shape[0], domain_shape[1]),
Stephan Seitz's avatar
Stephan Seitz committed
113
114
115
116
117
118
                             c=c_arr,
                             v=v_arr,
                             c_next=c_arr)

    # write down advection diffusion pde
    # the equation is represented by a single term and an implicit "=0" is assumed.
Stephan Seitz's avatar
Stephan Seitz committed
119
120
    adv_diff_pde = ps.fd.transient(c) - ps.fd.diffusion(
        c, sp.Symbol("D")) + ps.fd.advection(c, v)
Stephan Seitz's avatar
Stephan Seitz committed
121
122
123

    discretize = ps.fd.Discretization2ndOrder(args.dx, args.dt)
    discretization = discretize(adv_diff_pde)
Stephan Seitz's avatar
Stephan Seitz committed
124
125
    discretization = discretization.subs(sp.Symbol("D"),
                                         args.diffusion_coefficient)
Stephan Seitz's avatar
Stephan Seitz committed
126
127
128
129
    forward_assignments = ps.AssignmentCollection(
        [ps.Assignment(c_next.center(), discretization)], [])

    autodiff = pystencils_autodiff.AutoDiffOp(
Stephan Seitz's avatar
Stephan Seitz committed
130
131
        forward_assignments,
        diff_mode='transposed-forward')  # , constant_fields=[v]
Stephan Seitz's avatar
Stephan Seitz committed
132
133
134
135
136
137
138
139
140

    print('Forward assignments:')
    print(autodiff.forward_assignments)
    print('Backward assignments:')
    print(autodiff.backward_assignments)


def test_tfmad_vector_input_data():
    parser = argparse.ArgumentParser()
Stephan Seitz's avatar
Stephan Seitz committed
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
    parser.add_argument('--domain_shape',
                        default=(100, 30),
                        nargs=2,
                        type=int,
                        help="")
    parser.add_argument('--dx', default=1, type=float, help="")
    parser.add_argument('--dt', default=0.01, type=float, help="")
    parser.add_argument('--diffusion_coefficient',
                        default=1,
                        type=float,
                        help="")
    parser.add_argument('--num_total_time_steps', default=100, type=int)
    parser.add_argument('--num_time_steps_for_op', default=1, type=int)
    parser.add_argument('--learning_rate', default=1e-2, type=float)
    parser.add_argument('--dtype', default=np.float64, type=np.dtype)
    parser.add_argument('--num_optimization_steps', default=2000, type=int)
Stephan Seitz's avatar
Stephan Seitz committed
157
    parser.add_argument('vargs', nargs='*')
Stephan Seitz's avatar
Stephan Seitz committed
158
159
160
161
162
163

    args = parser.parse_args()
    check_tfmad_vector_input_data(args)


def test_tfmad_gradient_check_torch():
Stephan Seitz's avatar
Stephan Seitz committed
164
165
    torch = pytest.importorskip('torch')

Stephan Seitz's avatar
Stephan Seitz committed
166
167
168
169
170
171
172
173
174
175
176
177
178
    a, b, out = ps.fields("a, b, out: float[21,13]")

    cont = ps.fd.Diff(a, 0) - ps.fd.Diff(a, 1) - \
        ps.fd.Diff(b, 0) + ps.fd.Diff(b, 1)
    discretize = ps.fd.Discretization2ndOrder(dx=1)
    discretization = discretize(cont)

    assignment = ps.Assignment(out.center(), discretization)
    assignment_collection = ps.AssignmentCollection([assignment], [])
    print('Forward')
    print(assignment_collection)

    print('Backward')
Stephan Seitz's avatar
Stephan Seitz committed
179
180
    auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection,
                                               diff_mode='transposed-forward')
Stephan Seitz's avatar
Stephan Seitz committed
181
182
183
184
185
    backward = auto_diff.backward_assignments
    print(backward)
    print('Forward output fields (to check order)')
    print(auto_diff.forward_input_fields)

Stephan Seitz's avatar
Stephan Seitz committed
186
187
188
    a_tensor = torch.zeros(*a.shape, dtype=torch.float64, requires_grad=True)
    b_tensor = torch.zeros(*b.shape, dtype=torch.float64, requires_grad=True)

Stephan Seitz's avatar
Stephan Seitz committed
189
190
191
192
193
    function = auto_diff.create_tensorflow_op({
        a: a_tensor,
        b: b_tensor
    },
                                              backend='torch')
Stephan Seitz's avatar
Stephan Seitz committed
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240

    torch.autograd.gradcheck(function.apply, [a_tensor, b_tensor])


def get_curl(input_field: ps.Field, curl_field: ps.Field):
    """Return a ps.AssignmentCollection describing the calculation of
    the curl given a 2d or 3d vector field [z,y,x](f) or [y,x](f)

    Note that the curl of a 2d vector field is defined in ℝ3!
    Only the non-zero z-component is returned

    Arguments:
        field {ps.Field} -- A field with index_dimensions <= 1
            Scalar fields are interpreted as a z-component

    Raises:
        NotImplementedError -- [description]
        NotImplementedError -- Only support 2d or 3d vector fields or scalar fields are supported

    Returns:
        ps.AssignmentCollection -- AssignmentCollection describing the calculation of the curl
    """
    assert input_field.index_dimensions <= 1, "Must be a vector or a scalar field"
    assert curl_field.index_dimensions == 1, "Must be a vector field"
    discretize = ps.fd.Discretization2ndOrder(dx=1)

    if input_field.index_dimensions == 0:
        dy = ps.fd.Diff(input_field, 0)
        dx = ps.fd.Diff(input_field, 1)
        f_x = ps.Assignment(curl_field.center(0), discretize(dy))
        f_y = ps.Assignment(curl_field.center(1), discretize(dx))
        return ps.AssignmentCollection([f_x, f_y], [])

    else:

        if input_field.index_shape[0] == 2:
            raise NotImplementedError()

        elif input_field.index_shape[0] == 3:
            raise NotImplementedError()
        else:
            raise NotImplementedError()


def test_tfmad_two_outputs():

    domain_shape = (20, 30)
Stephan Seitz's avatar
Stephan Seitz committed
241
    vector_shape = domain_shape + (2, )
Stephan Seitz's avatar
Stephan Seitz committed
242

Stephan Seitz's avatar
Stephan Seitz committed
243
244
245
246
247
248
    curl_input_for_u = ps.Field.create_fixed_size(field_name='curl_input',
                                                  shape=domain_shape,
                                                  index_dimensions=0)
    u_field = ps.Field.create_fixed_size(field_name='curl',
                                         shape=vector_shape,
                                         index_dimensions=1)
Stephan Seitz's avatar
Stephan Seitz committed
249

Stephan Seitz's avatar
Stephan Seitz committed
250
251
252
    curl_op = pystencils_autodiff.AutoDiffOp(get_curl(curl_input_for_u,
                                                      u_field),
                                             diff_mode="transposed-forward")
Stephan Seitz's avatar
Stephan Seitz committed
253
254
255
256
257
258

    print('Forward')
    print(curl_op.forward_assignments)

    print('Backward')
    print(curl_op.backward_assignments)