Add array reflection.

1 files changed, 18 insertions, 0 deletions

diff --git a/utils.py b/utils.py
index b225bbe..7440919 100644
--- a/utils.py
+++ b/utils.py
@@ -19,3 +19,21 @@ def matrix_invert(mat):
     inv = np.linalg.inv(mat)
     return inv.transpose((1, 2, 0)).reshape(oldshape)
 
+def array_reflect(data, parity = 1.0, axis = -1):
+    """
+    Reflect an N-D array with respect to the specified axis. E.g. input
+    [0, 1, 2, 3] becomes [3, 2, 1, 0, 1, 2, 3] with parity=1.0 and
+    [-3, -2, -1, 0, 1, 2, 3] with parity=-1.0.
+
+    :param array_like data: The array to reflect.
+    :param float parity: The reflected portion is multiplied by this factor,
+                         typically 1.0 or -1.0.
+    :param int axis: Index of the axis to reflect along.
+    :return: Reflected array.
+    """
+    slices0 = [slice(None) for _ in data.shape]
+    slices1 = [slice(None) for _ in data.shape]
+    slices0[axis] = slice(None, None, -1)
+    slices1[axis] = slice(1, None)
+
+    return np.concatenate((parity * data[slices0], data[slices1]), axis)