diff --git a/src/pystencils_autodiff/__init__.py b/src/pystencils_autodiff/__init__.py index 289f2acfa6882824c82d1f589bb9a329f3f657ca..bc139b97fe2c981d40fe4800ff6a85e83f912019 100644 --- a/src/pystencils_autodiff/__init__.py +++ b/src/pystencils_autodiff/__init__.py @@ -1,12 +1,10 @@ import pystencils_autodiff.backends -# from pystencils_autodiff._field_to_tensors import ( -# tf_constant_from_field, tf_placeholder_from_field, -# tf_scalar_variable_from_field, tf_variable_from_field, -# torch_tensor_from_field) +from pystencils_autodiff._field_to_tensors import ( + tf_constant_from_field, tf_placeholder_from_field, tf_scalar_variable_from_field, + tf_variable_from_field, torch_tensor_from_field) from pystencils_autodiff.adjoint_field import AdjointField -from pystencils_autodiff.autodiff import (AutoDiffAstPair, AutoDiffOp, - create_backward_assignments, - get_jacobian_of_assignments) +from pystencils_autodiff.autodiff import ( + AutoDiffAstPair, AutoDiffOp, create_backward_assignments, get_jacobian_of_assignments) __all__ = ['backends', 'AdjointField', @@ -14,7 +12,6 @@ __all__ = ['backends', 'create_backward_assignments', 'AutoDiffOp', 'AutoDiffAstPair', - # "tf_constant_from_field", " tf_placeholder_from_field", - # "tf_scalar_variable_from_field", " tf_variable_from_field", - # "torch_tensor_from_field" - ] + "tf_constant_from_field", " tf_placeholder_from_field", + "tf_scalar_variable_from_field", " tf_variable_from_field", + "torch_tensor_from_field"] diff --git a/src/pystencils_autodiff/_field_to_tensors.py b/src/pystencils_autodiff/_field_to_tensors.py new file mode 100644 index 0000000000000000000000000000000000000000..0f4a7fec2c0ccc8407ee244af728641b98f424a6 --- /dev/null +++ b/src/pystencils_autodiff/_field_to_tensors.py @@ -0,0 +1,37 @@ +import numpy as np + +try: + import tensorflow as tf +except ImportError: + pass +try: + import torch +except ImportError: + pass + + +def tf_constant_from_field(field, init_val=0): + return tf.constant(init_val, dtype=field.dtype.numpy_dtype, shape=field.shape, name=field.name + '_constant') + + +def tf_scalar_variable_from_field(field, init_val, constraint=None): + var = tf.Variable(init_val, dtype=field.dtype.numpy_dtype, name=field.name + '_variable', constraint=constraint) + return var * tf_constant_from_field(field, 1) + + +def tf_variable_from_field(field, init_val=0, constraint=None): + if isinstance(init_val, (int, float)): + init_val *= np.ones(field.shape, field.dtype.numpy_dtype) + + return tf.Variable(init_val, dtype=field.dtype.numpy_dtype, name=field.name + '_variable', constraint=constraint) + + +def tf_placeholder_from_field(field): + return tf.placeholder(dtype=field.dtype.numpy_dtype, name=field.name + '_placeholder', shape=field.shape) + + +def torch_tensor_from_field(field, init_val=0, cuda=True, requires_grad=False): + if isinstance(init_val, (int, float)): + init_val *= np.ones(field.shape, field.dtype.numpy_dtype) + device = torch.device('cuda' if cuda else 'cpu') + return torch.tensor(init_val, requires_grad=requires_grad, device=device)