merged returns template and can be used with cut

test/test_core.py

-import tfields
 import numpy as np
-import unittest
 from tempfile import NamedTemporaryFile
 import pickle
+import unittest
+import tfields
 
 ATOL = 1e-8
 
@@ -39,16 +39,15 @@ class Base_Check(object):
 
     def test_basic_merge(self):
         # create 3 copies with different coord_sys
-        from IPython import embed; embed()
         merge_list = [self._inst.copy() for i in range(3)]
         merge_list[0].transform(tfields.bases.CARTESIAN)
         merge_list[1].transform(tfields.bases.CYLINDER)
         merge_list[2].transform(tfields.bases.SPHERICAL)
-        
+
         # merge them and check that the first coord_sys is taken
         obj = type(self._inst).merged(*merge_list)
         self.assertTrue(obj.coord_sys == tfields.bases.CARTESIAN)
-        
+
         # check that all copies are the same also with new coord_sys
         for i in range(len(merge_list)):
             value = np.allclose(merge_list[0],
@@ -114,7 +113,7 @@ class TensorMaps_Empty_Test(TensorFields_Empty_Test):
         self._inst = tfields.TensorMaps([], dim=3)
         self._maps = []
         self._maps_fields = []
-    
+
 
 class TensorFields_Copy_Test(TensorFields_Empty_Test):
     def setUp(self):

tfields/core.py

@@ -655,12 +655,19 @@ class Tensors(AbstractNdarray):
         Factory method
         Merges all tensor inputs to one tensor
 
+        Args:
+            **kwargs: passed to cls
+            dim (int):
+            return_templates (bool): return the templates which can be used
+                together with cut to retrieve the original objects
+
         Examples:
             >>> import numpy as np
             >>> import tfields
             >>> import tfields.bases
 
-            Use of most frequent coordinate system
+            The new object with turn out in the most frequent coordinate
+            system if not specified explicitly
 
             >>> vec_a = tfields.Tensors([[0, 0, 0], [0, 0, 1], [0, -1, 0]])
             >>> vec_b = tfields.Tensors([[5, 4, 1]], coord_sys=tfields.bases.cylinder)
@@ -686,18 +693,29 @@ class Tensors(AbstractNdarray):
             ...                                          len(merge) + 3,
             ...                                          1))])
 
-            >>> obj_list = [tfields.Tensors([[1, 2, 3]], coord_sys=tfields.bases.CYLINDER),
+            >>> obj_list = [tfields.Tensors([[1, 2, 3]],
+            ...             coord_sys=tfields.bases.CYLINDER),
             ...             tfields.Tensors([[3] * 3]),
             ...             tfields.Tensors([[5, 1, 3]])]
-            >>> merge2 = tfields.Tensors.merged(*obj_list, coord_sys=tfields.bases.CARTESIAN)
+            >>> merge2 = tfields.Tensors.merged(
+            ...     *obj_list, coord_sys=tfields.bases.CARTESIAN)
             >>> assert merge2.equal([[-0.41614684, 0.90929743, 3.],
             ...                      [3, 3, 3], [5, 1, 3]], atol=1e-8)
 
+            The return_templates argument allows to retrieve a template which
+            can be used with the cut method.
+
+            >>> merge, templates = tfields.Tensors.merged(
+            ...     vec_a, vec_b, vec_c, return_templates=True)
+            >>> assert merge.cut(templates[0]).equal(vec_a)
+            >>> assert merge.cut(templates[1]).equal(vec_b)
+            >>> assert merge.cut(templates[2]).equal(vec_c)
+
         """
 
         """ get most frequent coord_sys or predefined coord_sys """
         coord_sys = kwargs.get("coord_sys", None)
-        dimension = kwargs.get("dim", None)
+        return_templates = kwargs.pop("return_templates", False)
         if coord_sys is None:
             bases = []
             for t in objects:
@@ -735,11 +753,26 @@ class Tensors(AbstractNdarray):
         for i, obj in enumerate(remainingObjects):
             tensors = np.append(tensors, obj, axis=0)
 
-        if len(tensors) == 0 and dimension is None:
+        if len(tensors) == 0 and not 'dim' in kwargs:
+            # if you can not determine the tensor dimension, search for the
+            # first object with some entries
             for obj in objects:
-                kwargs["dim"] = dim(obj)
+                if len(obj) != 0:
+                    kwargs['dim'] = dim(obj)
+                    break
 
-        return cls.__new__(cls, tensors, **kwargs)
+        if not return_templates:
+            return cls.__new__(cls, tensors, **kwargs)
+        else:
+            tensor_lengths = [len(o) for o in objects]
+            cum_tensor_lengths = [sum(tensor_lengths[:i])
+                                  for i in range(len(objects))]
+            templates = [
+                tfields.TensorFields(
+                    obj,
+                    np.arange(tensor_lengths[i]) + cum_tensor_lengths[i])
+                for i, obj in enumerate(objects)]
+            return cls.__new__(cls, tensors, **kwargs), templates
 
     @classmethod
     def grid(cls, *base_vectors, **kwargs):
@@ -1241,14 +1274,36 @@ class Tensors(AbstractNdarray):
             mask = tfields.evalf(np.array(self), expression, coords=coords)
         return mask
 
-    def cut(self, expression, coord_sys=None):
+    def _cut_sympy(self, expression):
+        if len(self) == 0:
+            return self.copy()
+        mask = self.evalf(expression)  # coord_sys is handled by tmp_transform
+        mask.astype(bool)
+        inst = self[mask].copy()
+
+        # template
+        indices = np.arange(len(self))[mask]
+        # TODO: maybe np.empty instead of inst below. Reason why i copy is
+        # that the template is such marked for where it comes from.
+        # We have a trade-off with small memory consumpiton here.
+        template = tfields.TensorFields(inst, indices)
+        return inst, template
+
+    def _cut_template(self, template):
+        return self[template.fields[0]]
+
+    def cut(self, expression, coord_sys=None, return_template=False, **kwargs):
         """
-        Default cut method for Points3D. Works on a copy.
+        Extract a part of the object according to the logic given
+        by <expression>.
 
         Args:
-            expression (sympy logical expression): logical expression which will be evalfuated.
-                             use symbols x, y and z
-            coord_sys (str): coord_sys to evalfuate the expression in.
+            expression (sympy logical expression|tfields.TensorFields): logical
+                expression which will be evaluated. use symbols x, y and z.
+                If tfields.TensorFields or subclass is given, the expression
+                refers to a template.
+            coord_sys (str): coord_sys to evaluate the expression in. Only
+                active for template expression
 
         Examples:
             >>> import tfields
@@ -1264,19 +1319,44 @@ class Tensors(AbstractNdarray):
 
             combinations of cuts
 
-            >>> p.cut((x > 0) & (z < 0)).equal([[1, 2, -6], [1, 0, -1]])
+            >>> cut_expression = (x > 0) & (z < 0)
+            >>> combi_cut = p.cut(cut_expression)
+            >>> combi_cut.equal([[1, 2, -6], [1, 0, -1]])
+            True
+
+            Templates can be used to speed up the repeated cuts on the same
+            underlying tensor with the same expression but new fields.
+            First let us cut a but request the template on return:
+            >>> field1 = list(range(len(p)))
+            >>> tf = tfields.TensorFields(p, field1)
+            >>> tf_cut, template = tf.cut(cut_expression,
+            ...                           return_template=True)
+
+            Now repeat the cut with a new field:
+            >>> field2 = p
+            >>> tf.fields.append(field2)
+            >>> tf_template_cut = tf.cut(template)
+            >>> tf_template_cut.equal(combi_cut)
+            True
+            >>> tf_template_cut.fields[0].equal([2, 4])
+            True
+            >>> tf_template_cut.fields[1].equal(combi_cut)
             True
 
         Returns:
             copy of self with cut applied
+            [optional: template - requires <return_template> switch]
 
         """
-        if len(self) == 0:
-            return self.copy()
-        mask = self.evalf(expression, coord_sys=coord_sys or self.coord_sys)
-        mask.astype(bool)
-        inst = self[mask].copy()
-        return inst
+        with self.tmp_transform(coord_sys or self.coord_sys):
+            if issubclass(type(expression), TensorFields):
+                template = expression
+                obj = self._cut_template(template)
+            else:
+                obj, template = self._cut_sympy(expression, **kwargs)
+        if return_template:
+            return obj, template
+        return obj
 
     def distances(self, other, **kwargs):
         """
@@ -1676,7 +1756,13 @@ class TensorFields(Tensors):
                 "Got objects of types {types} instead.".format(**locals())
             )
 
-        inst = super(TensorFields, cls).merged(*objects, **kwargs)
+        return_value = super(TensorFields, cls).merged(*objects, **kwargs)
+        return_templates = kwargs.get('return_templates', False)
+        if return_templates:
+            inst, templates = return_value
+        else:
+            inst, templates = (return_value, None)
+
 
         fields = []
         if all([len(obj.fields) == len(objects[0].fields) for obj in objects]):
@@ -1686,7 +1772,10 @@ class TensorFields(Tensors):
                 )
                 fields.append(field)
         inst = cls.__new__(cls, inst, *fields)
-        return inst
+        if return_templates:
+            return inst, templates
+        else:
+            return inst
 
     @property
     def names(self):
@@ -1830,6 +1919,7 @@ class TensorMaps(TensorFields):
     __slots__ = ['coord_sys', 'name', 'fields', 'maps']
 
     def __new__(cls, tensors, *fields, **kwargs):
+        # TODO: the type of maps should rather be a dictionary with keys: dim
         maps = kwargs.pop("maps", [])
         maps_cp = []
         for mp in maps:
@@ -1946,26 +2036,53 @@ class TensorMaps(TensorFields):
                 )
             )
         tensor_lengths = [len(o) for o in objects]
-        cum_tensor_lengths = [
-            sum(tensor_lengths[:i]) for i in range(len(objects))
-        ]
+        cum_tensor_lengths = [sum(tensor_lengths[:i])
+                              for i in range(len(objects))]
 
-        maps = []
-        dims = []
-        for i, o in enumerate(objects):
-            for map_field in o.maps:
+        return_value = super(TensorMaps, cls).merged(*objects, **kwargs)
+        return_templates = kwargs.get('return_templates', False)
+        if return_templates:
+            inst, templates = return_value
+        else:
+            inst, templates = (return_value, None)
+
+        dim_maps_dict = {}  # {dim: {(obj_index(i), map_index(j): maps_field}}
+        for i, obj in enumerate(objects):
+            for j, map_field in enumerate(obj.maps):
                 map_field = map_field + cum_tensor_lengths[i]
-                try:
-                    mp_idx = dims.index(map_field.dim)
-                except ValueError:
-                    maps.append(map_field)
-                    dims.append(map_field.dim)
-                else:
-                    maps[mp_idx] = TensorFields.merged(maps[mp_idx], map_field)
+                if map_field.dim not in dim_maps_dict:
+                    dim_maps_dict[map_field.dim] = {}
+                dim_maps_dict[map_field.dim][(i, j)] = map_field
+
+        # for i, obj in enumerate(objects):
+        maps = []
+        for dimension in sorted(dim_maps_dict.keys()):
+            # sort by object index
+            keys = dim_maps_dict[dimension].keys()
+            obj_indices = [key[0] for key in keys]
+            map_indices = [key[1] for key in keys]
+            dim_maps = [dim_maps_dict[dimension][key] for key in keys]
+
+            obj_indices, map_indices, dim_maps = tfields.lib.util.multi_sort(
+                obj_indices, map_indices, dim_maps)
+            return_value = TensorFields.merged(
+                *dim_maps,
+                return_templates=return_templates,
+            )
+            if return_templates:
+                mp, map_templates = return_value
+                for i in obj_indices:
+                    j = map_indices[i]
+                    templates[i].maps[j] = map_templates[i]
+            else:
+                mp = return_value
+            maps.append(mp)
 
-        inst = super(TensorMaps, cls).merged(*objects, **kwargs)
         inst = cls.__new__(cls, inst, maps=maps)
-        return inst
+        if 'return_templates' in kwargs:
+            return inst, templates
+        else:
+            return inst
 
     def equal(self, other, **kwargs):
         """

tfields/mesh3D.py

@@ -742,7 +742,7 @@ class Mesh3D(tfields.TensorMaps):
                                             **kwargs))
             if merge_functions is None:
                 merge_functions = [lambda x: np.mean(x, axis=0)] * len(fields)
-        
+
         # build point_face_assignment if not given.
         if point_face_assignment is not None:
             if len(point_face_assignment) != len(tensor_field):
@@ -764,7 +764,7 @@ class Mesh3D(tfields.TensorMaps):
         for field, map_field, merge_function in \
                 zip(fields, empty_map_fields, merge_functions):
             for i, f_index in enumerate(point_face_assignment_set):
-                # if i % 1000 == 0: 
+                # if i % 1000 == 0:
                 #     print(i, len(point_face_assignment_set))
                 if f_index == -1:
                     # point could not be mapped
@@ -1048,8 +1048,7 @@ class Mesh3D(tfields.TensorMaps):
                               maps=maps)
         return inst
 
-    def cut(self, expression, coord_sys=None, at_intersection=None,
-            return_template=False):
+    def cut(self, *args, **kwargs):
         """
         cut method for Mesh3D.
         Args:
@@ -1143,16 +1142,7 @@ class Mesh3D(tfields.TensorMaps):
             * optional: template
 
         """
-        with self.tmp_transform(coord_sys or self.coord_sys):
-            if isinstance(expression, Mesh3D):
-                template = expression
-                obj = self._cut_template(template)
-            else:
-                at_intersection = at_intersection or "remove"
-                obj, template = self._cut_sympy(expression, at_intersection=at_intersection)
-        if return_template:
-            return obj, template
-        return obj
+        return super().cut(*args, **kwargs)
 
     def disjoint_parts(self, return_template=False):
         """