<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst

from StringIO import StringIO

import numpy as np

from .. import core, funcs
from ...tests.helper import pytest
from ... import units as u

try:
    import scipy  # pylint: disable=W0611
except ImportError:
    HAS_SCIPY = False
else:
    HAS_SCIPY = True


def setup_function(function):
    # Make sure that tests don't affect default cosmology
    core.set_current('no_default')


def teardown_function(function):
    # Make sure that tests don't affect default cosmology
    core.set_current('no_default')


def test_basic():
    cosmo = core.FlatLambdaCDM(H0=70, Om0=0.27, Tcmb0=2.0, Neff=3.04)
    assert np.allclose(cosmo.Om0, 0.27)
    assert np.allclose(cosmo.Ode0, 0.729975, rtol=1e-4)
    assert np.allclose(cosmo.Ogamma0, 1.463285e-5, rtol=1e-4)
    assert np.allclose(cosmo.Onu0, 1.01026e-5, rtol=1e-4)
    assert np.allclose(cosmo.Ok0, 0.0)
    assert np.allclose(cosmo.Om0 + cosmo.Ode0 + cosmo.Ogamma0 + cosmo.Onu0,
                       1.0, rtol=1e-6)
    assert np.allclose(cosmo.Om(1) + cosmo.Ode(1) + cosmo.Ogamma(1) +
                       cosmo.Onu(1), 1.0, rtol=1e-6)
    assert np.allclose(cosmo.Tcmb0.value, 2.0)
    assert np.allclose(cosmo.Tnu0.value, 1.4275317, rtol=1e-5)
    assert np.allclose(cosmo.Neff, 3.04)
    assert np.allclose(cosmo.h, 0.7)
    assert np.allclose(cosmo.H0.value, 70.0)

    # Make sure setting them as quantities gives the same results
    H0 = u.Quantity(70, u.km / (u.s * u.Mpc))
    T = u.Quantity(2.0, u.K)
    cosmo = core.FlatLambdaCDM(H0=H0, Om0=0.27, Tcmb0=T, Neff=3.04)
    assert np.allclose(cosmo.Om0, 0.27)
    assert np.allclose(cosmo.Ode0, 0.729975, rtol=1e-4)
    assert np.allclose(cosmo.Ogamma0, 1.463285e-5, rtol=1e-4)
    assert np.allclose(cosmo.Onu0, 1.01026e-5, rtol=1e-4)
    assert np.allclose(cosmo.Ok0, 0.0)
    assert np.allclose(cosmo.Om0 + cosmo.Ode0 + cosmo.Ogamma0 + cosmo.Onu0,
                       1.0, rtol=1e-6)
    assert np.allclose(cosmo.Om(1) + cosmo.Ode(1) + cosmo.Ogamma(1) +
                       cosmo.Onu(1), 1.0, rtol=1e-6)
    assert np.allclose(cosmo.Tcmb0.value, 2.0)
    assert np.allclose(cosmo.Tnu0.value, 1.4275317, rtol=1e-5)
    assert np.allclose(cosmo.Neff, 3.04)
    assert np.allclose(cosmo.h, 0.7)
    assert np.allclose(cosmo.H0.value, 70.0)


@pytest.mark.skipif('not HAS_SCIPY')
def test_units():
    """ Test if the right units are being returned"""

    cosmo = core.FlatLambdaCDM(H0=70, Om0=0.27, Tcmb0=2.0)
    assert cosmo.comoving_distance(1.0).unit == u.Mpc
    assert cosmo.angular_diameter_distance(1.0).unit == u.Mpc
    assert cosmo.angular_diameter_distance_z1z2(1.0, 2.0).unit == u.Mpc
    assert cosmo.comoving_distance(1.0).unit == u.Mpc
    assert cosmo.luminosity_distance(1.0).unit == u.Mpc
    assert cosmo.lookback_time(1.0).unit == u.Gyr
    assert cosmo.H0.unit == u.km / u.Mpc / u.s
    assert cosmo.H(1.0).unit == u.km / u.Mpc / u.s
    assert cosmo.Tcmb0.unit == u.K
    assert cosmo.Tcmb(1.0).unit == u.K
    assert cosmo.Tcmb([0.0, 1.0]).unit == u.K
    assert cosmo.Tnu0.unit == u.K
    assert cosmo.Tnu(1.0).unit == u.K
    assert cosmo.Tnu([0.0, 1.0]).unit == u.K
    assert cosmo.arcsec_per_kpc_comoving(1.0).unit == u.arcsec / u.kpc
    assert cosmo.arcsec_per_kpc_proper(1.0).unit == u.arcsec / u.kpc
    assert cosmo.kpc_comoving_per_arcmin(1.0).unit == u.kpc / u.arcmin
    assert cosmo.kpc_proper_per_arcmin(1.0).unit == u.kpc / u.arcmin
    assert cosmo.critical_density(1.0).unit == u.g / u.cm ** 3
    assert cosmo.comoving_volume(1.0).unit == u.Mpc ** 3
    assert cosmo.age(1.0).unit == u.Gyr
    assert cosmo.distmod(1.0).unit == u.mag

def test_repr():
    """ Test string representation of built in classes"""
    cosmo = core.LambdaCDM(70, 0.3, 0.5)
    expected = 'LambdaCDM(H0=70 km / (Mpc s), Om0=0.3, Ode0=0.5, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(cosmo) == expected

    cosmo = core.LambdaCDM(70, 0.3, 0.5, m_nu=u.Quantity(0.01, u.eV))
    expected = 'LambdaCDM(H0=70 km / (Mpc s), Om0=0.3, Ode0=0.5, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.01  0.01  0.01] eV)'
    assert "{0:s}".format(cosmo) == expected

    cosmo = core.FlatLambdaCDM(50.0, 0.27)
    expected = 'FlatLambdaCDM(H0=50 km / (Mpc s), Om0=0.27, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(cosmo) == expected

    cosmo = core.wCDM(60.0, 0.27, 0.6, w0=-0.8, name='test1')
    expected = 'wCDM(name="test1", H0=60 km / (Mpc s), Om0=0.27, Ode0=0.6, w0=-0.8, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(cosmo) == expected

    cosmo = core.FlatwCDM(65.0, 0.27, w0=-0.6, name='test2')
    expected = 'FlatwCDM(name="test2", H0=65 km / (Mpc s), Om0=0.27, w0=-0.6, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(cosmo) == expected

    cosmo = core.w0waCDM(60.0, 0.25, 0.4, w0=-0.6, wa=0.1, name='test3')
    expected = 'w0waCDM(name="test3", H0=60 km / (Mpc s), Om0=0.25, Ode0=0.4, w0=-0.6, wa=0.1, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(cosmo) == expected

    cosmo = core.Flatw0waCDM(55.0, 0.35, w0=-0.9, wa=-0.2, name='test4')
    expected = 'Flatw0waCDM(name="test4", H0=55 km / (Mpc s), Om0=0.35, w0=-0.9, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(cosmo) == expected

    cosmo = core.wpwaCDM(50.0, 0.3, 0.3, wp=-0.9, wa=-0.2, zp=0.3, name='test5')
    expected = 'wpwaCDM(name="test5", H0=50 km / (Mpc s), Om0=0.3, Ode0=0.3, wp=-0.9, wa=-0.2, zp=0.3, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(cosmo) == expected

    cosmo = core.w0wzCDM(55.0, 0.4, 0.8, w0=-1.05, wz=-0.2,
                         m_nu=u.Quantity([0.001, 0.01, 0.015], u.eV))
    expected = 'w0wzCDM(H0=55 km / (Mpc s), Om0=0.4, Ode0=0.8, w0=-1.05, wz=-0.2 Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.001  0.01   0.015] eV)'
    assert "{0:s}".format(cosmo) == expected

@pytest.mark.skipif('not HAS_SCIPY')
def test_flat_z1():
    """ Test a flat cosmology at z=1 against several other on-line
    calculators.
    """
    cosmo = core.FlatLambdaCDM(H0=70, Om0=0.27, Tcmb0=0.0)
    z = 1

    # Test values were taken from the following web cosmology
    # calculators on 27th Feb 2012:

    # Wright: http://www.astro.ucla.edu/~wright/CosmoCalc.html
    #         (http://adsabs.harvard.edu/abs/2006PASP..118.1711W)
    # Kempner: http://www.kempner.net/cosmic.php
    # iCosmos: http://www.icosmos.co.uk/index.html

    # The order of values below is Wright, Kempner, iCosmos'
    assert np.allclose(cosmo.comoving_distance(z).value,
                       [3364.5, 3364.8, 3364.7988], rtol=1e-4)
    assert np.allclose(cosmo.angular_diameter_distance(z).value,
                       [1682.3, 1682.4, 1682.3994], rtol=1e-4)
    assert np.allclose(cosmo.luminosity_distance(z).value,
                       [6729.2, 6729.6, 6729.5976], rtol=1e-4)
    assert np.allclose(cosmo.lookback_time(z).value,
                       [7.841, 7.84178, 7.843], rtol=1e-3)


# This class is to test whether the routines work correctly
# if one only overloads w(z)<commit_msg>Fixed cosmology test for Python 3.4<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst

from StringIO import StringIO

import numpy as np

from .. import core, funcs
from ...tests.helper import pytest
from ... import units as u

try:
    import scipy  # pylint: disable=W0611
except ImportError:
    HAS_SCIPY = False
else:
    HAS_SCIPY = True


def setup_function(function):
    # Make sure that tests don't affect default cosmology
    core.set_current('no_default')


def teardown_function(function):
    # Make sure that tests don't affect default cosmology
    core.set_current('no_default')


def test_basic():
    cosmo = core.FlatLambdaCDM(H0=70, Om0=0.27, Tcmb0=2.0, Neff=3.04)
    assert np.allclose(cosmo.Om0, 0.27)
    assert np.allclose(cosmo.Ode0, 0.729975, rtol=1e-4)
    assert np.allclose(cosmo.Ogamma0, 1.463285e-5, rtol=1e-4)
    assert np.allclose(cosmo.Onu0, 1.01026e-5, rtol=1e-4)
    assert np.allclose(cosmo.Ok0, 0.0)
    assert np.allclose(cosmo.Om0 + cosmo.Ode0 + cosmo.Ogamma0 + cosmo.Onu0,
                       1.0, rtol=1e-6)
    assert np.allclose(cosmo.Om(1) + cosmo.Ode(1) + cosmo.Ogamma(1) +
                       cosmo.Onu(1), 1.0, rtol=1e-6)
    assert np.allclose(cosmo.Tcmb0.value, 2.0)
    assert np.allclose(cosmo.Tnu0.value, 1.4275317, rtol=1e-5)
    assert np.allclose(cosmo.Neff, 3.04)
    assert np.allclose(cosmo.h, 0.7)
    assert np.allclose(cosmo.H0.value, 70.0)

    # Make sure setting them as quantities gives the same results
    H0 = u.Quantity(70, u.km / (u.s * u.Mpc))
    T = u.Quantity(2.0, u.K)
    cosmo = core.FlatLambdaCDM(H0=H0, Om0=0.27, Tcmb0=T, Neff=3.04)
    assert np.allclose(cosmo.Om0, 0.27)
    assert np.allclose(cosmo.Ode0, 0.729975, rtol=1e-4)
    assert np.allclose(cosmo.Ogamma0, 1.463285e-5, rtol=1e-4)
    assert np.allclose(cosmo.Onu0, 1.01026e-5, rtol=1e-4)
    assert np.allclose(cosmo.Ok0, 0.0)
    assert np.allclose(cosmo.Om0 + cosmo.Ode0 + cosmo.Ogamma0 + cosmo.Onu0,
                       1.0, rtol=1e-6)
    assert np.allclose(cosmo.Om(1) + cosmo.Ode(1) + cosmo.Ogamma(1) +
                       cosmo.Onu(1), 1.0, rtol=1e-6)
    assert np.allclose(cosmo.Tcmb0.value, 2.0)
    assert np.allclose(cosmo.Tnu0.value, 1.4275317, rtol=1e-5)
    assert np.allclose(cosmo.Neff, 3.04)
    assert np.allclose(cosmo.h, 0.7)
    assert np.allclose(cosmo.H0.value, 70.0)


@pytest.mark.skipif('not HAS_SCIPY')
def test_units():
    """ Test if the right units are being returned"""

    cosmo = core.FlatLambdaCDM(H0=70, Om0=0.27, Tcmb0=2.0)
    assert cosmo.comoving_distance(1.0).unit == u.Mpc
    assert cosmo.angular_diameter_distance(1.0).unit == u.Mpc
    assert cosmo.angular_diameter_distance_z1z2(1.0, 2.0).unit == u.Mpc
    assert cosmo.comoving_distance(1.0).unit == u.Mpc
    assert cosmo.luminosity_distance(1.0).unit == u.Mpc
    assert cosmo.lookback_time(1.0).unit == u.Gyr
    assert cosmo.H0.unit == u.km / u.Mpc / u.s
    assert cosmo.H(1.0).unit == u.km / u.Mpc / u.s
    assert cosmo.Tcmb0.unit == u.K
    assert cosmo.Tcmb(1.0).unit == u.K
    assert cosmo.Tcmb([0.0, 1.0]).unit == u.K
    assert cosmo.Tnu0.unit == u.K
    assert cosmo.Tnu(1.0).unit == u.K
    assert cosmo.Tnu([0.0, 1.0]).unit == u.K
    assert cosmo.arcsec_per_kpc_comoving(1.0).unit == u.arcsec / u.kpc
    assert cosmo.arcsec_per_kpc_proper(1.0).unit == u.arcsec / u.kpc
    assert cosmo.kpc_comoving_per_arcmin(1.0).unit == u.kpc / u.arcmin
    assert cosmo.kpc_proper_per_arcmin(1.0).unit == u.kpc / u.arcmin
    assert cosmo.critical_density(1.0).unit == u.g / u.cm ** 3
    assert cosmo.comoving_volume(1.0).unit == u.Mpc ** 3
    assert cosmo.age(1.0).unit == u.Gyr
    assert cosmo.distmod(1.0).unit == u.mag

def test_repr():
    """ Test string representation of built in classes"""
    cosmo = core.LambdaCDM(70, 0.3, 0.5)
    expected = 'LambdaCDM(H0=70 km / (Mpc s), Om0=0.3, Ode0=0.5, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(str(cosmo)) == expected

    cosmo = core.LambdaCDM(70, 0.3, 0.5, m_nu=u.Quantity(0.01, u.eV))
    expected = 'LambdaCDM(H0=70 km / (Mpc s), Om0=0.3, Ode0=0.5, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.01  0.01  0.01] eV)'
    assert "{0:s}".format(str(cosmo)) == expected

    cosmo = core.FlatLambdaCDM(50.0, 0.27)
    expected = 'FlatLambdaCDM(H0=50 km / (Mpc s), Om0=0.27, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(str(cosmo)) == expected

    cosmo = core.wCDM(60.0, 0.27, 0.6, w0=-0.8, name='test1')
    expected = 'wCDM(name="test1", H0=60 km / (Mpc s), Om0=0.27, Ode0=0.6, w0=-0.8, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(str(cosmo)) == expected

    cosmo = core.FlatwCDM(65.0, 0.27, w0=-0.6, name='test2')
    expected = 'FlatwCDM(name="test2", H0=65 km / (Mpc s), Om0=0.27, w0=-0.6, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(str(cosmo)) == expected

    cosmo = core.w0waCDM(60.0, 0.25, 0.4, w0=-0.6, wa=0.1, name='test3')
    expected = 'w0waCDM(name="test3", H0=60 km / (Mpc s), Om0=0.25, Ode0=0.4, w0=-0.6, wa=0.1, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(str(cosmo)) == expected

    cosmo = core.Flatw0waCDM(55.0, 0.35, w0=-0.9, wa=-0.2, name='test4')
    expected = 'Flatw0waCDM(name="test4", H0=55 km / (Mpc s), Om0=0.35, w0=-0.9, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(str(cosmo)) == expected

    cosmo = core.wpwaCDM(50.0, 0.3, 0.3, wp=-0.9, wa=-0.2, zp=0.3, name='test5')
    expected = 'wpwaCDM(name="test5", H0=50 km / (Mpc s), Om0=0.3, Ode0=0.3, wp=-0.9, wa=-0.2, zp=0.3, Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.  0.  0.] eV)'
    assert "{0:s}".format(str(cosmo)) == expected

    cosmo = core.w0wzCDM(55.0, 0.4, 0.8, w0=-1.05, wz=-0.2,
                         m_nu=u.Quantity([0.001, 0.01, 0.015], u.eV))
    expected = 'w0wzCDM(H0=55 km / (Mpc s), Om0=0.4, Ode0=0.8, w0=-1.05, wz=-0.2 Tcmb0=2.725 K, Neff=3.04, m_nu=[ 0.001  0.01   0.015] eV)'
    assert "{0:s}".format(str(cosmo)) == expected

@pytest.mark.skipif('not HAS_SCIPY')
def test_flat_z1():
    """ Test a flat cosmology at z=1 against several other on-line
    calculators.
    """
    cosmo = core.FlatLambdaCDM(H0=70, Om0=0.27, Tcmb0=0.0)
    z = 1

    # Test values were taken from the following web cosmology
    # calculators on 27th Feb 2012:

    # Wright: http://www.astro.ucla.edu/~wright/CosmoCalc.html
    #         (http://adsabs.harvard.edu/abs/2006PASP..118.1711W)
    # Kempner: http://www.kempner.net/cosmic.php
    # iCosmos: http://www.icosmos.co.uk/index.html

    # The order of values below is Wright, Kempner, iCosmos'
    assert np.allclose(cosmo.comoving_distance(z).value,
                       [3364.5, 3364.8, 3364.7988], rtol=1e-4)
    assert np.allclose(cosmo.angular_diameter_distance(z).value,
                       [1682.3, 1682.4, 1682.3994], rtol=1e-4)
    assert np.allclose(cosmo.luminosity_distance(z).value,
                       [6729.2, 6729.6, 6729.5976], rtol=1e-4)
    assert np.allclose(cosmo.lookback_time(z).value,
                       [7.841, 7.84178, 7.843], rtol=1e-3)


# This class is to test whether the routines work correctly
# if one only overloads w(z)
{'fraction_non_alphanumeric': 0.1768296912937912, 'fraction_numerical': 0.09684356573014222, 'mean_word_length': 2.8391478029294275, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 6, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '356428', 'n_tokens_mistral': 7034, 'n_tokens_neox': 6446, 'n_words': 1234}