def df_fn(s, mobi_fn):
    """ Derivative (element-wise) of water fractional flow

    Parameters
    ----------
    s : ndarray, shape (ny, nx) | (ny*nx,)
        Saturation

    mobi_fn : callable
        Mobility function lamb_w, lamb_o, dlamb_w, dlamb_o = mobi_fn(s, deriv=True) where:
            lamb_w : water mobility
            lamb_o : oil mobility
            dlamb_w : derivative of water mobility
            dlamb_o : derivative of oil mobility
    """
    lamb_w, lamb_o, dlamb_w, dlamb_o = mobi_fn(s, deriv=True)
    return dlamb_w / (lamb_w + lamb_o) - lamb_w * (dlamb_w + dlamb_o) / (lamb_w + lamb_o)**2