def create_macroscopic_setter_op(self, velocity_input_tensor, density_input_tensor, force_input_tensor=None, backend='tensorflow', **kernel_compilation_kwargs): rho_field = self._data_handling.fields[self.density_data_name] vel_field = self._data_handling.fields[self.velocity_data_name] if self.force_data_name: assert force_input_tensor is not None, "must set force_field_tensor if using a force model" force_field = self._data_handling.fields[self.force_data_name] else: force_field = None _, setter_eqs = self._create_macroscopic_setter_and_getter_equations() op = pystencils_autodiff.AutoDiffOp( setter_eqs, self.name + "_SetMacroscopicValues", **kernel_compilation_kwargs) rtn_dict = op.create_tensorflow_op( { rho_field: density_input_tensor, vel_field: velocity_input_tensor, force_field: force_input_tensor }, backend=backend) return rtn_dict
def test_tfmad_gradient_check_torch(): torch = pytest.importorskip('torch') a, b, out = ps.fields("a, b, out: float[5,7]") cont = 2 * ps.fd.Diff(a, 0) - 1.5 * ps.fd.Diff(a, 1) \ - ps.fd.Diff(b, 0) + 3 * ps.fd.Diff(b, 1) discretize = ps.fd.Discretization2ndOrder(dx=1) discretization = discretize(cont) + 1.2 * a.center assignment = ps.Assignment(out.center(), discretization) assignment_collection = ps.AssignmentCollection([assignment], []) print('Forward') print(assignment_collection) print('Backward') auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection, diff_mode='transposed-forward') backward = auto_diff.backward_assignments print(backward) print('Forward output fields (to check order)') print(auto_diff.forward_input_fields) a_tensor = torch.zeros(*a.shape, dtype=torch.float64, requires_grad=True) b_tensor = torch.zeros(*b.shape, dtype=torch.float64, requires_grad=True) function = auto_diff.create_tensorflow_op({ a: a_tensor, b: b_tensor }, backend='torch') torch.autograd.gradcheck(function.apply, [a_tensor, b_tensor])
def test_valid_boundary_handling_tensorflow_native(): pytest.importorskip('tensorflow') import tensorflow as tf a, b, out = ps.fields("a, b, out: double[10,11]") print(a.shape) cont = 2 * ps.fd.Diff(a, 0) - 1.5 * ps.fd.Diff(a, 1) - ps.fd.Diff( b, 0) + 3 * ps.fd.Diff(b, 1) discretize = ps.fd.Discretization2ndOrder(dx=1) discretization = discretize(cont) assignment = ps.Assignment(out.center(), discretization + 0.1 * b[1, 0]) assignment_collection = ps.AssignmentCollection([assignment], []) print('Forward') print(assignment_collection) print('Backward') auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection, boundary_handling='valid') 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(use_cuda=with_cuda, backend='tensorflow_native') # print(out_tensor) out_tensor = auto_diff.create_tensorflow_op(use_cuda=False, backend='tensorflow_native')( a=a_tensor, b=b_tensor) with tf.compat.v1.Session() as sess: sess.run(tf.global_variables_initializer()) sess.run(out_tensor) if True: gradient_error = compute_gradient_error_without_border( [a_tensor, b_tensor], [a.shape, b.shape], out_tensor, out.shape, num_border_pixels=0, ndim=2, debug=False) print('error: %s' % gradient_error.max_error) print('avg error: %s' % gradient_error.avg_error) assert any(e < 1e-4 for e in gradient_error.values())
def test_tfmad_gradient_check_torch_native(with_offsets, with_cuda): torch = pytest.importorskip('torch') import torch a, b, out = ps.fields("a, b, out: float64[5,7]") if with_offsets: cont = 2 * ps.fd.Diff(a, 0) - 1.5 * ps.fd.Diff(a, 1) - ps.fd.Diff( b, 0) + 3 * ps.fd.Diff(b, 1) discretize = ps.fd.Discretization2ndOrder(dx=1) discretization = discretize(cont) assignment = ps.Assignment(out.center(), discretization + 1.2 * a.center()) else: assignment = ps.Assignment(out.center(), 1.2 * a.center + 0.1 * b.center) assignment_collection = ps.AssignmentCollection([assignment], []) print('Forward') print(assignment_collection) print('Backward') auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection, boundary_handling='zeros', diff_mode='transposed-forward') backward = auto_diff.backward_assignments print(backward) print('Forward output fields (to check order)') print(auto_diff.forward_input_fields) a_tensor = torch.zeros(*a.shape, dtype=torch.float64, requires_grad=True).contiguous() b_tensor = torch.zeros(*b.shape, dtype=torch.float64, requires_grad=True).contiguous() if with_cuda: a_tensor = a_tensor.cuda() b_tensor = b_tensor.cuda() function = auto_diff.create_tensorflow_op(use_cuda=with_cuda, backend='torch_native') dict = {a: a_tensor, b: b_tensor} torch.autograd.gradcheck( function.apply, tuple([dict[f] for f in auto_diff.forward_input_fields]), atol=1e-4, raise_exception=True)
def test_tfmad_gradient_check_tensorflow_native(with_offsets, with_cuda, gradient_check): pytest.importorskip('tensorflow') import tensorflow as tf a, b, out = ps.fields("a, b, out: double[21,13]", layout='fzyx') print(a.shape) if with_offsets: cont = 2 * ps.fd.Diff(a, 0) - 1.5 * ps.fd.Diff(a, 1) - ps.fd.Diff( b, 0) + 3 * ps.fd.Diff(b, 1) discretize = ps.fd.Discretization2ndOrder(dx=1) discretization = discretize(cont) assignment = ps.Assignment(out.center(), discretization + 0.1 * b[1, 0]) else: assignment = ps.Assignment(out.center(), 1.2 * a.center + 0.1 * b.center) assignment_collection = ps.AssignmentCollection([assignment], []) print('Forward') print(assignment_collection) print('Backward') auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection, boundary_handling='zeros') backward = auto_diff.backward_assignments print(backward) print('Forward output fields (to check order)') print(auto_diff.forward_input_fields) # out_tensor = auto_diff.create_tensorflow_op(use_cuda=with_cuda, backend='tensorflow_native') # print(out_tensor) a_tensor = tf.Variable(np.zeros(a.shape, a.dtype.numpy_dtype)) b_tensor = tf.Variable(np.zeros(a.shape, a.dtype.numpy_dtype)) op = auto_diff.create_tensorflow_op(use_cuda=with_cuda, backend='tensorflow_native') theoretical, numerical = tf.test.compute_gradient(op, [a_tensor, b_tensor], delta=0.001) assert np.allclose(theoretical[0], numerical[0]) assert np.allclose(theoretical[1], numerical[1])
def test_tfmad_two_outputs(): domain_shape = (20, 30) vector_shape = domain_shape + (2, ) 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) curl_op = pystencils_autodiff.AutoDiffOp(get_curl(curl_input_for_u, u_field), diff_mode="transposed-forward") print('Forward') print(curl_op.forward_assignments) print('Backward') print(curl_op.backward_assignments)
def test_tfmad_gradient_check(): tf = pytest.importorskip('tensorflow') a, b, out = ps.fields("a, b, out: double[5,6]") print(a.shape) 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') auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection, diff_mode='transposed-forward') 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( [a_tensor, b_tensor], [a.shape, b.shape], out_tensor, out.shape, num_border_pixels=2, ndim=2) print('error: %s' % gradient_error.max_error) assert any(e < 1e-4 for e in gradient_error.values())
def create_macroscopic_getter_op(self, pdf_input_tensor, force_field_tensor=None, backend='tensorflow', **kernel_compilation_kwargs): """Creates a Tensorflow Op or a operation for another framework Arguments: pdf_input_tensor {[type]} -- [description] Keyword Arguments: force_field_tensor {[type]} -- [description] (default: {None}) backend {str} -- [description] (default: {'tensorflow'}) Returns: [type] -- [description] """ pdf_field = self._data_handling.fields[self._pdf_arr_name] if self.force_data_name: assert force_field_tensor is not None, "must set force_field_tensor if using a force model" force_field = self._data_handling.fields[self.force_data_name] else: force_field = None getter_eqs, _ = self._create_macroscopic_setter_and_getter_equations() op = pystencils_autodiff.AutoDiffOp( getter_eqs, self.name + "_SetMacroscopicValues", **kernel_compilation_kwargs) return op.create_tensorflow_op( { pdf_field: pdf_input_tensor, force_field: force_field_tensor }, backend=backend)
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') auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection, diff_mode='transposed-forward') backward = auto_diff.backward_assignments print(backward) print('Forward output fields (to check order)') print(auto_diff.forward_input_fields) print(auto_diff)
def test_tfmad_gradient_check_two_outputs(with_cuda): torch = pytest.importorskip('torch') import torch a, b, out1, out2, out3 = ps.fields( "a, b, out1, out2, out3: float64[21,13]") assignment_collection = ps.AssignmentCollection({ out1.center: a.center + b.center, out2.center: a.center - b.center, out3.center: sympy.exp(b[-1, 0]) }) print('Forward') print(assignment_collection) print('Backward') auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection, boundary_handling='zeros', diff_mode='transposed-forward') print(auto_diff.backward_fields) backward = auto_diff.backward_assignments print(backward) print('Forward output fields (to check order)') print(auto_diff.forward_input_fields) a_tensor = torch.zeros(*a.shape, dtype=torch.float64, requires_grad=True).contiguous() b_tensor = torch.zeros(*b.shape, dtype=torch.float64, requires_grad=True).contiguous() out1_tensor = torch.zeros(*a.shape, dtype=torch.float64, requires_grad=True).contiguous() out2_tensor = torch.zeros(*b.shape, dtype=torch.float64, requires_grad=True).contiguous() out3_tensor = torch.zeros(*b.shape, dtype=torch.float64, requires_grad=True).contiguous() if with_cuda: a_tensor = a_tensor.cuda() b_tensor = b_tensor.cuda() out1_tensor = out1_tensor.cuda() out2_tensor = out2_tensor.cuda() out3_tensor = out3_tensor.cuda() function = auto_diff.create_tensorflow_op(use_cuda=with_cuda, backend='torch_native') dict = { a: a_tensor, b: b_tensor, out1: out1_tensor, out2: out2_tensor, out3: out3_tensor, } torch.autograd.gradcheck( function.apply, tuple([dict[f] for f in auto_diff.forward_input_fields]), atol=1e-4, raise_exception=True)
def __init__(self, update_rule: LbmCollisionRule, src_pdf_field: Union[Field, str] = '', tmp_pdf_field: Union[Field, str] = '', time_constant_fields=[], *args, constant_fields=[], optimization={}, additional_fields=[], **method_parameters): self._adjoint_boundary_conditions = dict() src_pdf_field, tmp_pdf_field = _guess_src_dst_field_from_update_rule( update_rule, src_pdf_field, tmp_pdf_field) self.pdf_field = src_pdf_field self.temporary_field = tmp_pdf_field # Field.layout is just spatial try: real_layout = pystencils.field.get_layout_from_strides( src_pdf_field.strides) ndim = len(src_pdf_field.shape) if pystencils.field.layout_string_to_tuple( 'fzyx', dim=ndim) == real_layout: layout_string = 'fzyx' elif pystencils.field.layout_string_to_tuple( 'tensorflow', dim=ndim) == real_layout: layout_string = 'tensorflow' elif pystencils.field.layout_string_to_tuple( 'zyxf', dim=ndim) == real_layout: layout_string = 'zyxf' elif pystencils.field.layout_string_to_tuple( 'c', dim=ndim) == real_layout: layout_string = 'c' else: raise NotImplementedError( "Cannot recognize layout of src_pdf_field") optimization.update({'field_layout': layout_string}) except Exception: pass optimization.update({ 'double_precision': src_pdf_field.dtype.numpy_dtype == np.float64 }) super().__init__(*args, update_rule=update_rule, optimization=optimization, pdf_arr_name=src_pdf_field.name, tmp_arr_name=tmp_pdf_field.name, **method_parameters) self._target = 'gpu' if self._gpu else 'cpu' self._forward_assignments = update_rule # TODO: optimization parameters, e.g. target self._autodiff = pystencils_autodiff.AutoDiffOp( self._forward_assignments, "LBM", ghost_layers=self._data_handling.default_ghost_layers, target=self._target, data_type="double" if src_pdf_field.dtype.numpy_dtype == np.float64 else "float", cpu_openmp=True, time_constant_fields=time_constant_fields, constant_fields=constant_fields) self._lbmKernels = [ self._autodiff.get_forward_kernel(is_gpu=self._gpu) ] self._backwardLbmKernels = [ self._autodiff.get_backward_kernel(is_gpu=self._gpu) ] for field in set(additional_fields).union(set(time_constant_fields)): if field.name not in self.data_handling.cpu_arrays: self._data_handling.add_array(field.name, field.values_per_cell, layout=field.layout, dtype=field.dtype.numpy_dtype, gpu=self._gpu) adjoint_field = pystencils_autodiff.AdjointField(field) if adjoint_field.name not in self.data_handling.cpu_arrays: self._data_handling.add_array(adjoint_field.name, field.values_per_cell, layout=field.layout, dtype=field.dtype.numpy_dtype, gpu=self._gpu) for forward, backward in self._autodiff.backward_fields_map.items(): if backward.name not in self._data_handling.cpu_arrays: self._data_handling.add_array_like(backward.name, forward.name, gpu=self._gpu) self._backward_boundary_handling = LatticeBoltzmannBoundaryHandling( self.method, self._data_handling, self._backward_tmp_array_name, name="backward_boundary_handling", flag_interface=None, # TODO target=self._target, openmp=True)