cf-python/cf/test/test_HEALPix_utils.py at 9c28b796b00c0c53d6c18240bce676e802802170 · sadielbartholomew/cf-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
import datetime
import unittest

import numpy as np

import cf

healpix_imported = True
try:
    import healpix  # noqa: F401
except ImportError:
    healpix_imported = False


class DataTest(unittest.TestCase):
    """Unit tests for HEALPix utilities."""

    def setUp(self):
        """Preparations called immediately before each test method."""
        # Skip all if healpix module not available!
        if not healpix_imported:
            self.skipTest(
                "Test module requires 'healpix' package. Install it to run all."
            )

        # Create matching lists of selected nested, ring, nuniq and zuniq
        # indices for every refinement level.
        indices = [
            (r, i, healpix.nest2ring(healpix.order2nside(r), i))
            for r in range(30)
            for i in (0, 7, (12 * 4**r) - 1)
        ]
        self.refinement_levels, self.nested_indices, self.ring_indices = map(
            list, zip(*indices)
        )

        self.nuniq_indices = [  # noqa: F841
            i + 4 ** (1 + r)
            for r, i in zip(self.refinement_levels, self.nested_indices)
        ]

        self.zuniq_indices = [  # noqa: F841
            (2 * i + 1) * 4 ** (29 - r)
            for r, i in zip(self.refinement_levels, self.nested_indices)
        ]

    def test_HEALPix_uniq2zuniq(self):
        """Test _uniq2zuniq"""
        from cf.data.dask_utils_healpix import _uniq2zuniq

        self.assertTrue(
            np.array_equal(_uniq2zuniq(self.nuniq_indices), self.zuniq_indices)
        )

    def test_HEALPix_zuniq2uniq(self):
        """Test _zuniq2uniq"""
        from cf.data.dask_utils_healpix import _zuniq2uniq

        self.assertTrue(
            np.array_equal(_zuniq2uniq(self.zuniq_indices), self.nuniq_indices)
        )

    def test_HEALPix_zuniq2pix(self):
        """Test _zuniq2pix"""
        from cf.data.dask_utils_healpix import _zuniq2pix

        # nested
        order, i = _zuniq2pix(self.zuniq_indices, nest=True)

        self.assertTrue(np.array_equal(order, self.refinement_levels))
        self.assertTrue(np.array_equal(i, self.nested_indices))

        # ring
        with self.assertRaises(NotImplementedError):
            _zuniq2pix(self.zuniq_indices, nest=False)

    def test_HEALPix_pix2zuniq(self):
        """Test _pix2zuniq"""
        from cf.data.dask_utils_healpix import _pix2zuniq

        # nested
        z = [
            _pix2zuniq(r, i, nest=True)
            for r, i in zip(self.refinement_levels, self.nested_indices)
        ]

        self.assertTrue(np.array_equal(z, self.zuniq_indices))

        # ring
        z = [
            _pix2zuniq(r, i, nest=False)
            for r, i in zip(self.refinement_levels, self.ring_indices)
        ]

        self.assertTrue(np.array_equal(z, self.zuniq_indices))


if __name__ == "__main__":
    print("Run date:", datetime.datetime.now())
    cf.environment()
    print()
    unittest.main(verbosity=2)