Repository: scipy/scipy
File: scipy/integrate/_quadrature.py (Python)

Code before:
    nd = len(y.shape)
    if start is None:
        start = 0
    step = 2
    slice_all = (slice(None),)*nd
    slice0 = tupleset(slice_all, axis, slice(start, stop, step))
    slice1 = tupleset(slice_all, axis, slice(start+1, stop+1, step))
    slice2 = tupleset(slice_all, axis, slice(start+2, stop+2, step))
    slice_early = tupleset(slice_all, axis, slice(None, -1))
    slice_late = tupleset(slice_all, axis, slice(1, None))

    if x is None:  # Even-spaced Simpson's rule.
        result = xp.sum(y[slice0] + 4.0*y[slice1] + y[slice2], axis=axis)
        result *= dx / 3.0
    else:
        # Account for possibly different spacings.
        #    Simpson's rule changes a bit.
        h = x[slice_late] - x[slice_early]
        sl0 = tupleset(slice_all, axis, slice(start, stop, step))
        sl1 = tupleset(slice_all, axis, slice(start+1, stop+1, step))
        h0 = xp.astype(h[sl0], xp.float64, copy=False)
        h1 = xp.astype(h[sl1], xp.float64, copy=False)
        hsum = h0 + h1
        hprod = h0 * h1
        with np.errstate(divide='ignore', invalid='ignore'):
            zero = xp.zeros_like(h0)
            h0divh1 = xp.where(h1 == 0, zero, h0 / h1)
            h1divh0 = xp.where(h0 == 0, zero, h1 / h0)
            hsumdivhprod = xp.where(hprod == 0, zero, hsum / hprod)
            tmp = hsum/6.0 * (y[slice0] * (2.0 - h1divh0) +
                              y[slice1] * (hsum * hsumdivhprod) +
                              y[slice2] * (2.0 - h0divh1))
        result = xp.sum(tmp, axis=axis)
    return result


def simpson(y, *, x=None, dx=1.0, axis=-1):
    """
    Integrate y(x) using samples along the given axis and the composite
    Simpson's rule. If x is None, spacing of dx is assumed.

    Parameters
    ----------
    y : array_like
        Array to be integrated.
    x : array_like, optional
        If given, the points at which `y` is sampled.
    dx : float, optional
        Spacing of integration points along axis of `x`. Only used when
        `x` is None. Default is 1.
    axis : int, optional

Review comment (performance):
Is creating this zero array and the several `xp.where`s  going to lead to performance issues for large h0/h1/hprod/hsum? Is a lazywhere better, or even fixing up afterwards.

Code after:
        shape = list(res.shape)
        shape[axis] = 1
        res = xp.concat((xp.full(tuple(shape), initial, dtype=res.dtype), res),
                        axis=axis)

    return res


def _basic_simpson(y, start, stop, x, dx, axis, *, xp):
    nd = y.ndim
    if start is None:
        start = 0
    step = 2
    slice_all = (slice(None),)*nd
    slice0 = tupleset(slice_all, axis, slice(start, stop, step))
    slice1 = tupleset(slice_all, axis, slice(start+1, stop+1, step))
    slice2 = tupleset(slice_all, axis, slice(start+2, stop+2, step))

    if x is None:  # Even-spaced Simpson's rule.
        result = dx / 3.0 * xp.sum(y[slice0] + 4.0*y[slice1] + y[slice2], axis=axis)
    else:
        # Account for possibly different spacings.
        #    Simpson's rule changes a bit.
        h = xp.diff(x, axis=axis)
        sl0 = tupleset(slice_all, axis, slice(start, stop, step))
        sl1 = tupleset(slice_all, axis, slice(start+1, stop+1, step))
        h0 = h[sl0] if xp_size(h) else h
        h1 = h[sl1] if xp_size(h) else h
        hsum = h0 + h1
        hprod = h0 * h1
        h0divh1 = xpx.apply_where(h1 != 0, (h0, h1), xp.divide, fill_value=0.)
        tmp = hsum/6.0 * (y[slice0] *
            (2.0 - xpx.apply_where(h0divh1 != 0, (xp.ones_like(h0divh1), h0divh1,),
                                   xp.divide, fill_value=0.))
            + y[slice1] * (hsum * xpx.apply_where(hprod != 0, (hsum, hprod,),
                                                  xp.divide, fill_value=0.))
            + (y[slice2] if xp_size(y) > 1 else y[0:0, ...]) * (2.0 - h0divh1))
        result = xp.sum(tmp, axis=axis)
    return result


@xp_capabilities()
def simpson(y, x=None, *, dx=1.0, axis=-1):
    """
    Integrate y(x) using samples along the given axis and the composite
    Simpson's rule. If x is None, spacing of dx is assumed.

    Parameters
    ----------
    y : array_like
        Array to be integrated.