Add TypedMatrixSymbol (for usage of `MatrixSymbol` in kernels)

I don't know whether this is a good idea but SymPy supports assigning MatrixSymbols. Like

>>>A = MatrixSymbols('A', 3, 3)
>>>B = MatrixSymbols('B', 3, 3)
In [12]: pystencils.Assignment(A, B)                                                                                                 
Out[12]: A := B

With this hack I can generate code like this:

#define FUNC_PREFIX static

  2 FUNC_PREFIX void kernel(float * RESTRICT _data_y, int64_t const _size_y_0, int64_t const _size_y_1, int64_t const _size_y_2, int64_t
    const _stride_y_0, int64_t const _stride_y_1, int64_t const _stride_y_2, std::function< Vector3 < double >(int, int, int) > my_fun)
  3 {
  4    for (int ctr_0 = 0; ctr_0 < _size_y_0; ctr_0 += 1)
  5    {
  6       float * RESTRICT _data_y_00 = _data_y + _stride_y_0*ctr_0;
  7       for (int ctr_1 = 0; ctr_1 < _size_y_1; ctr_1 += 1)
  8       {
  9          float * RESTRICT _data_y_00_10 = _stride_y_1*ctr_1 + _data_y_00;
 10          for (int ctr_2 = 0; ctr_2 < _size_y_2; ctr_2 += 1)
 11          {
 12             const Vector3<double> A = my_fun(ctr_0, ctr_1, ctr_2);
 13             _data_y_00_10[_stride_y_2*ctr_2] = A[0] + A[1] + A[2];
 14          }
 15       }
 16    }
 17 }
  1 #define FUNC_PREFIX static
  2 template <class Functor_T>
  3 FUNC_PREFIX void kernel(float * RESTRICT _data_y, int64_t const _size_y_0, int64_t const _size_y_1, int64_t const _size_y_2, int64_t
    const _stride_y_0, int64_t const _stride_y_1, int64_t const _stride_y_2, Functor_T my_fun)
  4 {
  5    for (int ctr_0 = 0; ctr_0 < _size_y_0; ctr_0 += 1)
  6    {
  7       float * RESTRICT _data_y_00 = _data_y + _stride_y_0*ctr_0;
  8       for (int ctr_1 = 0; ctr_1 < _size_y_1; ctr_1 += 1)
  9       {
 10          float * RESTRICT _data_y_00_10 = _stride_y_1*ctr_1 + _data_y_00;
 11          for (int ctr_2 = 0; ctr_2 < _size_y_2; ctr_2 += 1)
 12          {
 13             const Vector3<double> A = my_fun(ctr_0, ctr_1, ctr_2);
 14             _data_y_00_10[_stride_y_2*ctr_2] = A[0] + A[1] + A[2];
 15          }
 16       }
 17    }
 18 }

from

    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())

pystencils/test_matrix_symbols.py

+import sympy as sp
+
+import pystencils
+from pystencils.data_types import TypedMatrixSymbol, TypedSymbol, create_type
+
+
+class DynamicFunction(sp.Function):
+    """
+    Function that is passed as an argument to a kernel.
+    Can be printed for example as `std::function` or as a functor template.
+    """
+
+    def __new__(cls, typed_function_symbol, return_dtype, *args):
+        obj = sp.Function.__new__(cls, typed_function_symbol, return_dtype, *args)
+        if hasattr(return_dtype, 'shape'):
+            obj.shape = return_dtype.shape
+
+        return obj
+
+    @property
+    def function_dtype(self):
+        return self.args[0].dtype
+
+    @property
+    def dtype(self):
+        return self.args[1].dtype
+
+    @property
+    def name(self):
+        return self.args[0].name
+
+    def __str__(self):
+        return f'{self.name}({", ".join(str(a) for a in self.args[2:])})'
+
+    def __repr__(self):
+        return self.__str__()
+
+
+def test_dynamic_matrix_location_dependent():
+    x, y = pystencils.fields('x, y:  float32[3d]')
+
+    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)