1# -*- coding: utf-8 -*-
2
3"""Testing :mod:`astropy.cosmology.units`."""
4
5##############################################################################
6# IMPORTS
7
8import contextlib
9
10import pytest
11
12import astropy.cosmology.units as cu
13import astropy.units as u
14from astropy.cosmology import Planck13, default_cosmology, flrw
15from astropy.tests.helper import assert_quantity_allclose
16from astropy.utils.compat.optional_deps import HAS_ASDF, HAS_SCIPY
17from astropy.utils.exceptions import AstropyDeprecationWarning
18
19##############################################################################
20# TESTS
21##############################################################################
22
23
24def test_has_expected_units():
25    """
26    Test that this module has the expected set of units. Some of the units are
27    imported from :mod:`astropy.units`, or vice versa. Here we test presence,
28    not usage. Units from :mod:`astropy.units` are tested in that module. Units
29    defined in :mod:`astropy.cosmology` will be tested subsequently.
30    """
31    with pytest.warns(AstropyDeprecationWarning, match="`littleh`"):
32        assert u.astrophys.littleh is cu.littleh
33
34
35def test_has_expected_equivalencies():
36    """
37    Test that this module has the expected set of equivalencies. Many of the
38    equivalencies are imported from :mod:`astropy.units`, so here we test
39    presence, not usage. Equivalencies from :mod:`astropy.units` are tested in
40    that module. Equivalencies defined in :mod:`astropy.cosmology` will be
41    tested subsequently.
42    """
43    with pytest.warns(AstropyDeprecationWarning, match="`with_H0`"):
44        assert u.equivalencies.with_H0 is cu.with_H0
45
46
47def test_littleh():
48    """Test :func:`astropy.cosmology.units.with_H0`."""
49    H0_70 = 70 * u.km / u.s / u.Mpc
50    h70dist = 70 * u.Mpc / cu.littleh
51
52    assert_quantity_allclose(h70dist.to(u.Mpc, cu.with_H0(H0_70)), 100 * u.Mpc)
53
54    # make sure using the default cosmology works
55    cosmodist = default_cosmology.get().H0.value * u.Mpc / cu.littleh
56    assert_quantity_allclose(cosmodist.to(u.Mpc, cu.with_H0()), 100 * u.Mpc)
57
58    # Now try a luminosity scaling
59    h1lum = 0.49 * u.Lsun * cu.littleh ** -2
60    assert_quantity_allclose(h1lum.to(u.Lsun, cu.with_H0(H0_70)), 1 * u.Lsun)
61
62    # And the trickiest one: magnitudes.  Using H0=10 here for the round numbers
63    H0_10 = 10 * u.km / u.s / u.Mpc
64    # assume the "true" magnitude M = 12.
65    # Then M - 5*log_10(h)  = M + 5 = 17
66    withlittlehmag = 17 * (u.mag - u.MagUnit(cu.littleh ** 2))
67    assert_quantity_allclose(withlittlehmag.to(u.mag, cu.with_H0(H0_10)), 12 * u.mag)
68
69
70@pytest.mark.skipif(not HAS_SCIPY, reason="Cosmology needs scipy")
71def test_dimensionless_redshift():
72    """Test :func:`astropy.cosmology.units.dimensionless_redshift`."""
73    z = 3 * cu.redshift
74    val = 3 * u.one
75
76    # show units not equal
77    assert z.unit == cu.redshift
78    assert z.unit != u.one
79
80    # test equivalency enabled by default
81    assert z == val
82
83    # also test that it works for powers
84    assert (3 * cu.redshift ** 3) == val
85
86    # and in composite units
87    assert (3 * u.km / cu.redshift ** 3) == 3 * u.km
88
89    # test it also works as an equivalency
90    with u.set_enabled_equivalencies([]):  # turn off default equivalencies
91        assert z.to(u.one, equivalencies=cu.dimensionless_redshift()) == val
92
93        with pytest.raises(ValueError):
94            z.to(u.one)
95
96    # if this fails, something is really wrong
97    with u.add_enabled_equivalencies(cu.dimensionless_redshift()):
98        assert z == val
99
100
101@pytest.mark.skipif(not HAS_SCIPY, reason="Cosmology needs scipy")
102def test_redshift_temperature():
103    """Test :func:`astropy.cosmology.units.redshift_temperature`."""
104    cosmo = Planck13.clone(Tcmb0=3 * u.K)
105    default_cosmo = default_cosmology.get()
106    z = 15 * cu.redshift
107    Tcmb = cosmo.Tcmb(z)
108
109    # 1) Default (without specifying the cosmology)
110    with default_cosmology.set(cosmo):
111        equivalency = cu.redshift_temperature()
112        assert_quantity_allclose(z.to(u.K, equivalency), Tcmb)
113        assert_quantity_allclose(Tcmb.to(cu.redshift, equivalency), z)
114
115    # showing the answer changes if the cosmology changes
116    # this test uses the default cosmology
117    equivalency = cu.redshift_temperature()
118    assert_quantity_allclose(z.to(u.K, equivalency), default_cosmo.Tcmb(z))
119    assert default_cosmo.Tcmb(z) != Tcmb
120
121    # 2) Specifying the cosmology
122    equivalency = cu.redshift_temperature(cosmo)
123    assert_quantity_allclose(z.to(u.K, equivalency), Tcmb)
124    assert_quantity_allclose(Tcmb.to(cu.redshift, equivalency), z)
125
126    # Test `atzkw`
127    equivalency = cu.redshift_temperature(cosmo, ztol=1e-10)
128    assert_quantity_allclose(Tcmb.to(cu.redshift, equivalency), z)
129
130
131@pytest.mark.skipif(not HAS_SCIPY, reason="Cosmology needs scipy")
132def test_redshift_hubble():
133    """Test :func:`astropy.cosmology.units.redshift_hubble`."""
134    unit = u.km / u.s / u.Mpc
135    cosmo = Planck13.clone(H0=100 * unit)
136    default_cosmo = default_cosmology.get()
137    z = 15 * cu.redshift
138    H = cosmo.H(z)
139    h = H.to_value(u.km/u.s/u.Mpc) / 100 * cu.littleh
140
141    # 1) Default (without specifying the cosmology)
142    with default_cosmology.set(cosmo):
143        equivalency = cu.redshift_hubble()
144        # H
145        assert_quantity_allclose(z.to(unit, equivalency), H)
146        assert_quantity_allclose(H.to(cu.redshift, equivalency), z)
147        # little-h
148        assert_quantity_allclose(z.to(cu.littleh, equivalency), h)
149        assert_quantity_allclose(h.to(cu.redshift, equivalency), z)
150
151    # showing the answer changes if the cosmology changes
152    # this test uses the default cosmology
153    equivalency = cu.redshift_hubble()
154    assert_quantity_allclose(z.to(unit, equivalency), default_cosmo.H(z))
155    assert default_cosmo.H(z) != H
156
157    # 2) Specifying the cosmology
158    equivalency = cu.redshift_hubble(cosmo)
159    # H
160    assert_quantity_allclose(z.to(unit, equivalency), H)
161    assert_quantity_allclose(H.to(cu.redshift, equivalency), z)
162    # little-h
163    assert_quantity_allclose(z.to(cu.littleh, equivalency), h)
164    assert_quantity_allclose(h.to(cu.redshift, equivalency), z)
165
166    # Test `atzkw`
167    equivalency = cu.redshift_hubble(cosmo, ztol=1e-10)
168    assert_quantity_allclose(H.to(cu.redshift, equivalency), z)  # H
169    assert_quantity_allclose(h.to(cu.redshift, equivalency), z)  # little-h
170
171
172@pytest.mark.skipif(not HAS_SCIPY, reason="Cosmology needs scipy")
173@pytest.mark.parametrize(
174    "kind",
175    [cu.redshift_distance.__defaults__[-1], "comoving", "lookback", "luminosity"]
176)
177def test_redshift_distance(kind):
178    """Test :func:`astropy.cosmology.units.redshift_distance`."""
179    z = 15 * cu.redshift
180    d = getattr(Planck13, kind + "_distance")(z)
181
182    equivalency = cu.redshift_distance(cosmology=Planck13, kind=kind)
183
184    # properties of Equivalency
185    assert equivalency.name[0] == "redshift_distance"
186    assert equivalency.kwargs[0]["cosmology"] == Planck13
187    assert equivalency.kwargs[0]["distance"] == kind
188
189    # roundtrip
190    assert_quantity_allclose(z.to(u.Mpc, equivalency), d)
191    assert_quantity_allclose(d.to(cu.redshift, equivalency), z)
192
193
194def test_redshift_distance_wrong_kind():
195    """Test :func:`astropy.cosmology.units.redshift_distance` wrong kind."""
196    with pytest.raises(ValueError, match="`kind`"):
197        cu.redshift_distance(kind=None)
198
199
200@pytest.mark.skipif(not HAS_SCIPY, reason="Cosmology needs scipy")
201class Test_with_redshift:
202    """Test `astropy.cosmology.units.with_redshift`."""
203
204    @pytest.fixture
205    def cosmo(self):
206        """Test cosmology."""
207        return Planck13.clone(Tcmb0=3 * u.K)
208
209    # ===========================================
210
211    def test_cosmo_different(self, cosmo):
212        """The default is different than the test cosmology."""
213        default_cosmo = default_cosmology.get()
214        assert default_cosmo != cosmo  # shows changing default
215
216    def test_no_equivalency(self, cosmo):
217        """Test the equivalency  ``with_redshift`` without any enabled."""
218        equivalency = cu.with_redshift(distance=None, hubble=False, Tcmb=False)
219        assert len(equivalency) == 0
220
221    # -------------------------------------------
222
223    def test_temperature_off(self, cosmo):
224        """Test ``with_redshift`` with the temperature off."""
225        default_cosmo = default_cosmology.get()
226        z = 15 * cu.redshift
227        Tcmb = cosmo.Tcmb(z)
228
229        # 1) Default (without specifying the cosmology)
230        with default_cosmology.set(cosmo):
231            equivalency = cu.with_redshift(Tcmb=False)
232            with pytest.raises(u.UnitConversionError, match="'redshift' and 'K'"):
233                z.to(u.K, equivalency)
234
235        # 2) Specifying the cosmology
236        equivalency = cu.with_redshift(cosmo, Tcmb=False)
237        with pytest.raises(u.UnitConversionError, match="'redshift' and 'K'"):
238            z.to(u.K, equivalency)
239
240    def test_temperature(self, cosmo):
241        """Test temperature equivalency component."""
242        default_cosmo = default_cosmology.get()
243        z = 15 * cu.redshift
244        Tcmb = cosmo.Tcmb(z)
245
246        # 1) Default (without specifying the cosmology)
247        with default_cosmology.set(cosmo):
248            equivalency = cu.with_redshift(Tcmb=True)
249            assert_quantity_allclose(z.to(u.K, equivalency), Tcmb)
250            assert_quantity_allclose(Tcmb.to(cu.redshift, equivalency), z)
251
252        # showing the answer changes if the cosmology changes
253        # this test uses the default cosmology
254        equivalency = cu.with_redshift(Tcmb=True)
255        assert_quantity_allclose(z.to(u.K, equivalency), default_cosmo.Tcmb(z))
256        assert default_cosmo.Tcmb(z) != Tcmb
257
258        # 2) Specifying the cosmology
259        equivalency = cu.with_redshift(cosmo, Tcmb=True)
260        assert_quantity_allclose(z.to(u.K, equivalency), Tcmb)
261        assert_quantity_allclose(Tcmb.to(cu.redshift, equivalency), z)
262
263        # Test `atzkw`
264        # this is really just a test that 'atzkw' doesn't fail
265        equivalency = cu.with_redshift(cosmo, Tcmb=True, atzkw={"ztol": 1e-10})
266        assert_quantity_allclose(Tcmb.to(cu.redshift, equivalency), z)
267
268    # -------------------------------------------
269
270    def test_hubble_off(self, cosmo):
271        """Test ``with_redshift`` with Hubble off."""
272        unit = u.km / u.s / u.Mpc
273        default_cosmo = default_cosmology.get()
274        z = 15 * cu.redshift
275        H = cosmo.H(z)
276
277        # 1) Default (without specifying the cosmology)
278        with default_cosmology.set(cosmo):
279            equivalency = cu.with_redshift(hubble=False)
280            with pytest.raises(u.UnitConversionError, match="'redshift' and 'km / "):
281                z.to(unit, equivalency)
282
283        # 2) Specifying the cosmology
284        equivalency = cu.with_redshift(cosmo, hubble=False)
285        with pytest.raises(u.UnitConversionError, match="'redshift' and 'km / "):
286            z.to(unit, equivalency)
287
288    def test_hubble(self, cosmo):
289        """Test Hubble equivalency component."""
290        unit = u.km/u.s/u.Mpc
291        default_cosmo = default_cosmology.get()
292        z = 15 * cu.redshift
293        H = cosmo.H(z)
294        h = H.to_value(u.km / u.s / u.Mpc) / 100 * cu.littleh
295
296        # 1) Default (without specifying the cosmology)
297        with default_cosmology.set(cosmo):
298            equivalency = cu.with_redshift(hubble=True)
299            # H
300            assert_quantity_allclose(z.to(unit, equivalency), H)
301            assert_quantity_allclose(H.to(cu.redshift, equivalency), z)
302            # little-h
303            assert_quantity_allclose(z.to(cu.littleh, equivalency), h)
304            assert_quantity_allclose(h.to(cu.redshift, equivalency), z)
305
306        # showing the answer changes if the cosmology changes
307        # this test uses the default cosmology
308        equivalency = cu.with_redshift(hubble=True)
309        assert_quantity_allclose(z.to(unit, equivalency), default_cosmo.H(z))
310        assert default_cosmo.H(z) != H
311
312        # 2) Specifying the cosmology
313        equivalency = cu.with_redshift(cosmo, hubble=True)
314        # H
315        assert_quantity_allclose(z.to(unit, equivalency), H)
316        assert_quantity_allclose(H.to(cu.redshift, equivalency), z)
317        # little-h
318        assert_quantity_allclose(z.to(cu.littleh, equivalency), h)
319        assert_quantity_allclose(h.to(cu.redshift, equivalency), z)
320
321        # Test `atzkw`
322        # this is really just a test that 'atzkw' doesn't fail
323        equivalency = cu.with_redshift(cosmo, hubble=True, atzkw={"ztol": 1e-10})
324        assert_quantity_allclose(H.to(cu.redshift, equivalency), z)  # H
325        assert_quantity_allclose(h.to(cu.redshift, equivalency), z)  # h
326
327    # -------------------------------------------
328
329    def test_distance_off(self, cosmo):
330        """Test ``with_redshift`` with the distance off."""
331        default_cosmo = default_cosmology.get()
332        z = 15 * cu.redshift
333
334        # 1) Default (without specifying the cosmology)
335        with default_cosmology.set(cosmo):
336            equivalency = cu.with_redshift(distance=None)
337            with pytest.raises(u.UnitConversionError, match="'redshift' and 'Mpc'"):
338                z.to(u.Mpc, equivalency)
339
340        # 2) Specifying the cosmology
341        equivalency = cu.with_redshift(cosmo, distance=None)
342        with pytest.raises(u.UnitConversionError, match="'redshift' and 'Mpc'"):
343            z.to(u.Mpc, equivalency)
344
345    def test_distance_default(self):
346        """Test distance equivalency default."""
347        z = 15 * cu.redshift
348        d = default_cosmology.get().comoving_distance(z)
349
350        equivalency = cu.with_redshift()
351        assert_quantity_allclose(z.to(u.Mpc, equivalency), d)
352        assert_quantity_allclose(d.to(cu.redshift, equivalency), z)
353
354    def test_distance_wrong_kind(self):
355        """Test distance equivalency, but the wrong kind."""
356        with pytest.raises(ValueError, match="`kind`"):
357            cu.with_redshift(distance=ValueError)
358
359    @pytest.mark.parametrize("kind", ["comoving", "lookback", "luminosity"])
360    def test_distance(self, kind):
361        """Test distance equivalency."""
362        cosmo = Planck13
363        z = 15 * cu.redshift
364        dist = getattr(cosmo, kind + "_distance")(z)
365
366        default_cosmo = default_cosmology.get()
367        assert default_cosmo != cosmo  # shows changing default
368
369        # 1) without specifying the cosmology
370        with default_cosmology.set(cosmo):
371            equivalency = cu.with_redshift(distance=kind)
372            assert_quantity_allclose(z.to(u.Mpc, equivalency), dist)
373
374        # showing the answer changes if the cosmology changes
375        # this test uses the default cosmology
376        equivalency = cu.with_redshift(distance=kind)
377        assert_quantity_allclose(z.to(u.Mpc, equivalency),
378                                 getattr(default_cosmo, kind + "_distance")(z))
379        assert not u.allclose(getattr(default_cosmo, kind + "_distance")(z), dist)
380
381        # 2) Specifying the cosmology
382        equivalency = cu.with_redshift(cosmo, distance=kind)
383        assert_quantity_allclose(z.to(u.Mpc, equivalency), dist)
384        assert_quantity_allclose(dist.to(cu.redshift, equivalency), z)
385
386        # Test atzkw
387        # this is really just a test that 'atzkw' doesn't fail
388        equivalency = cu.with_redshift(cosmo, distance=kind, atzkw={"ztol": 1e-10})
389        assert_quantity_allclose(dist.to(cu.redshift, equivalency), z)
390
391
392# FIXME! get "dimensionless_redshift", "with_redshift" to work in this
393# they are not in ``astropy.units.equivalencies``, so the following fails
394@pytest.mark.skipif(not HAS_ASDF, reason="requires ASDF")
395@pytest.mark.parametrize("equiv", [cu.with_H0])
396def test_equivalencies_asdf(tmpdir, equiv, recwarn):
397    from asdf.tests import helpers
398
399    tree = {"equiv": equiv()}
400    helpers.assert_roundtrip_tree(tree, tmpdir)
401
402
403def test_equivalency_context_manager():
404    base_registry = u.get_current_unit_registry()
405
406    # check starting with only the dimensionless_redshift equivalency.
407    assert len(base_registry.equivalencies) == 1
408    assert str(base_registry.equivalencies[0][0]) == "redshift"
409
410
411@pytest.mark.skipif(not HAS_SCIPY, reason = "Cosmology needs scipy")
412class Test_unsuitable_units:
413    """Ensure that unsuitable units are rejected"""
414
415    def test_scalar(self):
416        with pytest.raises(u.UnitConversionError):
417            Planck13.Om(0. * u.m)
418
419    @pytest.mark.parametrize('method',
420                             ['Om', 'Ob', 'Odm', 'Ode', 'Ogamma', 'Onu', 'Tcmb', 'Tnu',
421                              'scale_factor', 'angular_diameter_distance', 'luminosity_distance'])
422    def test_unit_exception(self, method):
423        with pytest.raises(u.UnitConversionError):
424            getattr(Planck13, method)([0., 1.] * u.m)
425
426    @pytest.mark.parametrize('method',
427                             ['_EdS_age', '_flat_age', '_dS_lookback_time', 'efunc', 'inv_efunc',
428                              'Ok'])
429    def test_LambdaCDM_unit_exception(self, method):
430        LambdaCDMTest = flrw.LambdaCDM(70.0, 0.3, 0.5)
431        with pytest.raises(u.UnitConversionError):
432            getattr(LambdaCDMTest, method)([0., 1.] * u.m)
433
434    @pytest.mark.parametrize('method', ['efunc', 'inv_efunc'])
435    def test_FlatLambdaCDM_unit_exception(self, method):
436        FlatLambdaCDMTest = flrw.FlatLambdaCDM(70.0, 0.3)
437        with pytest.raises(u.UnitConversionError):
438            getattr(FlatLambdaCDMTest, method)([0., 1.] * u.m)
439
440    @pytest.mark.parametrize('method', ['de_density_scale', 'efunc', 'inv_efunc'])
441    def test_wCDM_unit_exception(self, method):
442        wCDMTest = flrw.wCDM(60.0, 0.27, 0.6, w0 = -0.8)
443        with pytest.raises(u.UnitConversionError):
444            getattr(wCDMTest, method)([0., 1.] * u.m)
445
446    @pytest.mark.parametrize('method', ['efunc', 'inv_efunc'])
447    def test_FlatwCDM_unit_exception(self, method):
448        FlatwCDMTest = flrw.FlatwCDM(65.0, 0.27, w0 = -0.6)
449        with pytest.raises(u.UnitConversionError):
450            getattr(FlatwCDMTest, method)([0., 1.] * u.m)
451
452    @pytest.mark.parametrize('method', ['w', 'de_density_scale'])
453    def test_w0waCDM_unit_exception(self, method):
454        w0waCDMTest = flrw.w0waCDM(H0 = 70, Om0 = 0.27, Ode0 = 0.5, w0 = -1.2,
455                                   wa = -0.2)
456        with pytest.raises(u.UnitConversionError):
457            getattr(w0waCDMTest, method)([0., 1.] * u.m)
458
459    @pytest.mark.parametrize('method', ['w', 'de_density_scale'])
460    def test_wpwaCDM_unit_exception(self, method):
461        wpwaCDMTest = flrw.wpwaCDM(H0 = 70, Om0 = 0.27, Ode0 = 0.5, wp = -1.2,
462                                   wa = -0.2, zp = 0.9)
463        with pytest.raises(u.UnitConversionError):
464            getattr(wpwaCDMTest, method)([0., 1.] * u.m)
465
466    @pytest.mark.parametrize('method', ['w', 'de_density_scale'])
467    def test_w0wzCDM14_unit_exception(self, method):
468        w0wzCDMTest = flrw.w0wzCDM(H0 = 70, Om0 = 0.27, Ode0 = 0.5, w0 = -1.2,
469                                   wz = 0.1)
470        with pytest.raises(u.UnitConversionError):
471            getattr(w0wzCDMTest, method)([0., 1.] * u.m)
472