def f_fn(s, mobi_fn):
    """ Water fractional flow

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

    mobi_fn : callable
        Mobility function lamb_w, lamb_o = mobi_fn(s) where:
            lamb_w : water mobility
            lamb_o : oil mobility
    """
    lamb_w, lamb_o = mobi_fn(s)
    return lamb_w / (lamb_w + lamb_o)