Repository: scipy/scipy
File: scipy/signal/_waveforms.py (Python)

Code before:
        ramp.  ``width=0.5`` produces a triangle wave.
        If an array, causes wave shape to change over time, and must be the
        same length as `t`.

    Returns
    -------
    y : ndarray
        Output array containing the sawtooth waveform.

    Examples
    --------
    A 5 Hz waveform sampled at 500 Hz for 1 second:

    >>> import numpy as np
    >>> from scipy import signal
    >>> import matplotlib.pyplot as plt
    >>> t = np.linspace(0, 1, 500)
    >>> plt.plot(t, signal.sawtooth(2 * np.pi * 5 * t))

    """
    xp = array_namespace(t, width)
    t, w = xp_promote(t, width, broadcast=True, force_floating=True, xp=xp)
    y = xp.zeros(t.shape, dtype=t.dtype)

    # width must be between 0 and 1 inclusive
    mask1 = (w > 1) | (w < 0)
    y = xpx.at(y, mask1).set(xp.nan)

    # take t modulo 2*pi
    tmod = t % (2*xp.pi)

    # on the interval 0 to width*2*pi function is tmod / (pi*w) - 1
    mask2 = ~mask1 & (tmod < w*2*xp.pi)
    y = xpx.at(y, mask2).set(tmod[mask2]/(xp.pi*w[mask2]) - 1)

    # on the interval width*2*pi to 2*pi function is (pi*(w+1)-tmod) / (pi*(1-w))
    mask3 = ~(mask1 | mask2)
    y = xpx.at(y, mask3).set(
        (xp.pi*(w[mask3] + 1) - tmod[mask3])/(xp.pi*(1 - w[mask3]))
    )
    return y


def square(t, duty=0.5):
    """
    Return a periodic square-wave waveform.

    The square wave has a period ``2*pi``, has value +1 from 0 to
    ``2*pi*duty`` and -1 from ``2*pi*duty`` to ``2*pi``. `duty` must be in
    the interval ``[0,1]``.


Review comment (suggestion):
I wouldn't mention it, but if there are more of these, would it be more convenient than direct translation to use `where` to nan-out the offending `w`s and `xpx.apply_where` to create the result, propagating the NaNs? The implementation would become even more compact.

Code after:
        ramp.  ``width=0.5`` produces a triangle wave.
        If an array, causes wave shape to change over time, and must be the
        same length as `t`.

    Returns
    -------
    y : ndarray
        Output array containing the sawtooth waveform.

    Examples
    --------
    A 5 Hz waveform sampled at 500 Hz for 1 second:

    >>> import numpy as np
    >>> from scipy import signal
    >>> import matplotlib.pyplot as plt
    >>> t = np.linspace(0, 1, 500)
    >>> plt.plot(t, signal.sawtooth(2 * np.pi * 5 * t))

    """
    xp = array_namespace(t, width)
    t, w = xp_promote(t, width, broadcast=True, force_floating=True, xp=xp)
    y = xp.zeros_like(t)

    # width must be between 0 and 1 inclusive
    mask1 = (w > 1) | (w < 0)
    y = xpx.at(y, mask1).set(xp.nan)

    # take t modulo 2*pi
    tmod = t % (2*xp.pi)

    # on the interval 0 to width*2*pi function is tmod / (pi*w) - 1
    mask2 = ~mask1 & (tmod < w*2*xp.pi)
    y = xpx.at(y, mask2).set(tmod[mask2]/(xp.pi*w[mask2]) - 1)

    # on the interval width*2*pi to 2*pi function is (pi*(w+1)-tmod) / (pi*(1-w))
    mask3 = ~(mask1 | mask2)
    y = xpx.at(y, mask3).set(
        (xp.pi*(w[mask3] + 1) - tmod[mask3])/(xp.pi*(1 - w[mask3]))
    )
    return y


def square(t, duty=0.5):
    """
    Return a periodic square-wave waveform.

    The square wave has a period ``2*pi``, has value +1 from 0 to
    ``2*pi*duty`` and -1 from ``2*pi*duty`` to ``2*pi``. `duty` must be in
    the interval ``[0,1]``.