jax.numpy.unwrap#

jax.numpy.unwrap(p, discont=None, axis=-1, period=6.283185307179586)[source]#

Unwrap a periodic signal.

JAX implementation of numpy.unwrap().

Parameters:

  • p (ArrayLike) – input array

  • discont (ArrayLike | None) – the maximum allowable discontinuity in the sequence. The default is period / 2

  • axis (int) – the axis along which to unwrap; defaults to -1

  • period (ArrayLike) – the period of the signal, which defaults to \(2\pi\)

Returns:

An unwrapped copy of p.

Return type:

Array

Examples

Consider a situation in which you are making measurements of the position of a rotating disk via the x and y locations of some point on that disk. The underlying variable is an always-increating angle which we’ll generate this way, using degrees for ease of representation:

>>> rng = np.random.default_rng(0)
>>> theta = rng.integers(0, 90, size=(20,)).cumsum()
>>> theta
array([ 76, 133, 179, 203, 230, 233, 239, 240, 255, 328, 386, 468, 513,
       567, 654, 719, 775, 823, 873, 957])

Our observations of this angle are the x and y coordinates, given by the sine and cosine of this underlying angle:

>>> x, y = jnp.sin(jnp.deg2rad(theta)), jnp.cos(jnp.deg2rad(theta))

Now, say that given these x and y coordinates, we wish to recover the original angle theta. We might do this via the atan2() function:

>>> theta_out = jnp.rad2deg(jnp.atan2(x, y)).round()
>>> theta_out
Array([  76.,  133.,  179., -157., -130., -127., -121., -120., -105.,
        -32.,   26.,  108.,  153., -153.,  -66.,   -1.,   55.,  103.,
        153., -123.], dtype=float32)

The first few values match the input angle theta above, but after this the values are wrapped because the sin and cos observations obscure the phase information. The purpose of the unwrap() function is to recover the original signal from this wrapped view of it:

>>> jnp.unwrap(theta_out, period=360)
Array([ 76., 133., 179., 203., 230., 233., 239., 240., 255., 328., 386.,
       468., 513., 567., 654., 719., 775., 823., 873., 957.],      dtype=float32)

It does this by assuming that the true underlying sequence does not differ by more than discont (which defaults to period / 2) within a single step, and when it encounters a larger discontinuity it adds factors of the period to the data. For periodic signals that satisfy this assumption, unwrap() can recover the original phased signal.