Fix rst doctests in cosmology package