base_vectors implemented

tfields/lib/grid.py

@@ -14,7 +14,7 @@ def ensure_complex(*base_vectors):
     return base_vectors
 
 
-def to_base_vectors(*base_vectors):
+def _to_base_vectors(*base_vectors):
     """
     Transform tuples to arrays with np.mgrid
     Args:
@@ -99,7 +99,7 @@ def igrid(*base_vectors, **kwargs):
         obj = np.mgrid[base_vectors]
         obj = obj.reshape(obj.shape[0], -1).T
     else:
-        base_vectors = to_base_vectors(*base_vectors)
+        base_vectors = _to_base_vectors(*base_vectors)
 
         obj = np.empty(shape=(functools.reduce(lambda x, y: x * y,
                                                map(len, base_vectors)),
@@ -123,6 +123,37 @@ def igrid(*base_vectors, **kwargs):
     return obj
 
 
+def base_vectors(array):
+    """
+    describe the array in terms of base vectors
+    Inverse function of igrid
+
+    Examples:
+        >>> import tfields
+        >>> grid = tfields.igrid((3, 5, 5j))
+        >>> tfields.lib.grid.base_vectors(grid[:, 0])
+        (3.0, 5.0, 5j)
+        >>> grid2 = tfields.igrid((3, 5, 5j),
+        ...                       (1, 2, 2j))
+        >>> grid_circle = tfields.igrid(*tfields.lib.grid.base_vectors(grid2))
+        >>> assert tfields.Tensors(grid_circle).equal(grid2)
+
+    """
+    if len(array.shape) == 1:
+        values = sorted(set(array))
+        spacing = complex(0, len(values))
+        vmin = min(values)
+        vmax = max(values)
+        return (vmin, vmax, spacing)
+    elif len(array.shape) == 2:
+        bases = []
+        for i in range(array.shape[1]):
+            bases.append(base_vectors(array[:, i]))
+        return bases
+    else:
+        raise NotImplementedError("Description yet only till rank 1")
+
+
 def swap_columns(array, *index_tuples):
     """
     Args: