Merge pull request #16987 from scimax/master

ENH: Phase unwrapping generalized to arbitrary interval size

1 files changed, 53 insertions, 14 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 9ed7700d2..783d45c2f 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -1496,26 +1496,40 @@ def angle(z, deg=False):
     return a
 
 
-def _unwrap_dispatcher(p, discont=None, axis=None):
+def _unwrap_dispatcher(p, discont=None, axis=None, *, period=None):
     return (p,)
 
 
 @array_function_dispatch(_unwrap_dispatcher)
-def unwrap(p, discont=pi, axis=-1):
-    """
-    Unwrap by changing deltas between values to 2*pi complement.
+def unwrap(p, discont=None, axis=-1, *, period=2*pi):
+    r"""
+    Unwrap by taking the complement of large deltas with respect to the period.
+
+    This unwraps a signal `p` by changing elements which have an absolute
+    difference from their predecessor of more than ``max(discont, period/2)``
+    to their `period`-complementary values.
 
-    Unwrap radian phase `p` by changing absolute jumps greater than
-    `discont` to their 2*pi complement along the given axis.
+    For the default case where `period` is :math:`2\pi` and is `discont` is
+    :math:`\pi`, this unwraps a radian phase `p` such that adjacent differences
+    are never greater than :math:`\pi` by adding :math:`2k\pi` for some
+    integer :math:`k`.
 
     Parameters
     ----------
     p : array_like
         Input array.
     discont : float, optional
-        Maximum discontinuity between values, default is ``pi``.
+        Maximum discontinuity between values, default is ``period/2``. 
+        Values below ``period/2`` are treated as if they were ``period/2``.
+        To have an effect different from the default, `discont` should be
+        larger than ``period/2``.
     axis : int, optional
         Axis along which unwrap will operate, default is the last axis.
+    period: float, optional
+        Size of the range over which the input wraps. By default, it is
+        ``2 pi``.
+        
+        .. versionadded:: 1.21.0
 
     Returns
     -------
@@ -1528,9 +1542,9 @@ def unwrap(p, discont=pi, axis=-1):
 
     Notes
     -----
-    If the discontinuity in `p` is smaller than ``pi``, but larger than
-    `discont`, no unwrapping is done because taking the 2*pi complement
-    would only make the discontinuity larger.
+    If the discontinuity in `p` is smaller than ``period/2``, 
+    but larger than `discont`, no unwrapping is done because taking 
+    the complement would only make the discontinuity larger.
 
     Examples
     --------
@@ -1540,19 +1554,44 @@ def unwrap(p, discont=pi, axis=-1):
     array([ 0.        ,  0.78539816,  1.57079633,  5.49778714,  6.28318531]) # may vary
     >>> np.unwrap(phase)
     array([ 0.        ,  0.78539816,  1.57079633, -0.78539816,  0.        ]) # may vary
-
+    >>> np.unwrap([0, 1, 2, -1, 0], period=4)
+    array([0, 1, 2, 3, 4])
+    >>> np.unwrap([ 1, 2, 3, 4, 5, 6, 1, 2, 3], period=6)
+    array([1, 2, 3, 4, 5, 6, 7, 8, 9])
+    >>> np.unwrap([2, 3, 4, 5, 2, 3, 4, 5], period=4)
+    array([2, 3, 4, 5, 6, 7, 8, 9])
+    >>> phase_deg = np.mod(np.linspace(0 ,720, 19), 360) - 180
+    >>> np.unwrap(phase_deg, period=360)
+    array([-180., -140., -100.,  -60.,  -20.,   20.,   60.,  100.,  140.,
+            180.,  220.,  260.,  300.,  340.,  380.,  420.,  460.,  500.,
+            540.])
     """
     p = asarray(p)
     nd = p.ndim
     dd = diff(p, axis=axis)
+    if discont is None:
+        discont = period/2
     slice1 = [slice(None, None)]*nd     # full slices
     slice1[axis] = slice(1, None)
     slice1 = tuple(slice1)
-    ddmod = mod(dd + pi, 2*pi) - pi
-    _nx.copyto(ddmod, pi, where=(ddmod == -pi) & (dd > 0))
+    dtype = np.result_type(dd, period)
+    if _nx.issubdtype(dtype, _nx.integer):
+        interval_high, rem = divmod(period, 2) 
+        boundary_ambiguous = rem == 0
+    else:
+        interval_high = period / 2
+        boundary_ambiguous = True
+    interval_low = -interval_high
+    ddmod = mod(dd - interval_low, period) + interval_low
+    if boundary_ambiguous:
+        # for `mask = (abs(dd) == period/2)`, the above line made
+        # `ddmod[mask] == -period/2`. correct these such that
+        # `ddmod[mask] == sign(dd[mask])*period/2`.
+        _nx.copyto(ddmod, interval_high,
+                   where=(ddmod == interval_low) & (dd > 0))
     ph_correct = ddmod - dd
     _nx.copyto(ph_correct, 0, where=abs(dd) < discont)
-    up = array(p, copy=True, dtype='d')
+    up = array(p, copy=True, dtype=dtype)
     up[slice1] = p[slice1] + ph_correct.cumsum(axis)
     return up