Hybrid solution including interval boundaries, and keeping backward compatibility without boundary

1 files changed, 30 insertions, 7 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index c5541cc86..772f44fcf 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -1482,12 +1482,12 @@ def angle(z, deg=False):
     return a
 
 
-def _unwrap_dispatcher(p, discont=None, axis=None, *, interval_size=2*pi):
+def _unwrap_dispatcher(p, discont=None, axis=None, *, interval_size=2*pi, min_val=None, max_val=None):
     return (p,)
 
 
 @array_function_dispatch(_unwrap_dispatcher)
-def unwrap(p, discont=None, axis=-1, *, interval_size=2*pi):
+def unwrap(p, discont=None, axis=-1, *, interval_size=2*pi, min_val=None, max_val=None):
     """
     Unwrap by changing deltas between values to complement.
     
@@ -1502,12 +1502,25 @@ def unwrap(p, discont=None, axis=-1, *, interval_size=2*pi):
     ----------
     p : array_like
         Input array.
-    interval_size: float, optional
-        size of the range over which the input wraps.
     discont : float, optional
         Maximum discontinuity between values, default is ``interval_size/2``.
     axis : int, optional
         Axis along which unwrap will operate, default is the last axis.
+    interval_size: float, optional
+        Size of the range over which the input wraps. By default, it is 2 pi.
+        If ``min_val`` and ``max_val`` are given, ``interval_size`` is ignored
+        and the interval size is ``max_val - min_val``. 
+    min_val, max_val: float, optional
+        Boundaries of the interval over which the input array is expected to 
+        wrap. By default, they are ``None`` and the interval is considered as 
+        ``[-interval_size, interval_size]``. In case the first value of the 
+        phase input array, ``p[0]``, is outside of the interval 
+        ``[min_val, max_val]`` it will be corrected by an integral multiple of 
+        the interval size such that it will be within the 
+        boundaries.
+        Both boundaries require each other. If only one boundary is 
+        provided without the other, it will be ignored. 
+    
     Returns
     -------
     out : ndarray
@@ -1540,11 +1553,20 @@ def unwrap(p, discont=None, axis=-1, *, interval_size=2*pi):
         180.,  220.,  260.,  300.,  340.,  380.,  420.,  460.,  500.,
         540.])
     """
-    if discont is None:
-        discont = interval_size/2
     p = asarray(p)
     nd = p.ndim
     dd = diff(p, axis=axis)
+    offset = 0
+    if (not min_val is None) and (not max_val is None):
+        interval_size = max_val - min_val
+        slice0list = [slice(None, None)]*nd     # full slices
+        slice0list[axis] = 0
+        slice0 = tuple(slice0list)
+        offset_mul = (p[slice0] - min_val)//interval_size
+        slice0list[axis] = None
+        offset = -offset_mul[tuple(slice0list)]*interval_size
+    if discont is None:
+        discont = interval_size/2
     slice1 = [slice(None, None)]*nd     # full slices
     slice1[axis] = slice(1, None)
     slice1 = tuple(slice1)
@@ -1552,8 +1574,9 @@ def unwrap(p, discont=None, axis=-1, *, interval_size=2*pi):
     _nx.copyto(ddmod, interval_size/2, where=(ddmod == -interval_size/2) & (dd > 0))
     ph_correct = ddmod - dd
     _nx.copyto(ph_correct, 0, where=abs(dd) < discont)
+    p += offset
     up = array(p, copy=True, dtype='d')
-    up[slice1] = p[slice1] + ph_correct.cumsum(axis)
+    up[slice1] = p[slice1] + ph_correct.cumsum(axis) 
     return up