1"""Tests for ``jplephem``.
2
3See the accompanying ``jpltest`` module for a more intense numerical
4test suite that can verify that ``jplephem`` delivers, over hundreds of
5examples, the same results as when the ephemerides are run at JPL.  This
6smaller and more feature-oriented suite can be run with::
7
8    python -m unittest discover jplephem
9
10"""
11import mmap
12import numpy as np
13import sys
14import tempfile
15from doctest import DocTestSuite, ELLIPSIS
16from functools import partial
17from io import BytesIO
18from jplephem import Ephemeris, commandline
19from jplephem.exceptions import OutOfRangeError
20from jplephem.daf import DAF, FTPSTR, NAIF_DAF
21from jplephem.spk import SPK
22from struct import Struct
23try:
24    from unittest import SkipTest, TestCase
25except ImportError:
26    from unittest2 import SkipTest, TestCase
27
28epsilon_m = 0.01
29target_names = {
30    'mercury barycenter': 1, # BARYCENTER w.r.t. 0 SOLAR SYSTEM BARYCENTER
31    'venus barycenter': 2, # BARYCENTER w.r.t. 0 SOLAR SYSTEM BARYCENTER
32    'earthmoon': 3,        # BARYCENTER w.r.t. 0 SOLAR SYSTEM BARYCENTER
33    'mars barycenter': 4,  # BARYCENTER w.r.t. 0 SOLAR SYSTEM BARYCENTER
34    'jupiter': 5,          # BARYCENTER w.r.t. 0 SOLAR SYSTEM BARYCENTER
35    'saturn': 6,           # BARYCENTER w.r.t. 0 SOLAR SYSTEM BARYCENTER
36    'uranus': 7,           # BARYCENTER w.r.t. 0 SOLAR SYSTEM BARYCENTER
37    'neptune': 8,          # BARYCENTER w.r.t. 0 SOLAR SYSTEM BARYCENTER
38    'pluto': 9,            # BARYCENTER w.r.t. 0 SOLAR SYSTEM BARYCENTER
39    'sun': 10,             # w.r.t. 0 SOLAR SYSTEM BARYCENTER
40    'mercury': 199,        # w.r.t. 1 MERCURY BARYCENTER
41    'venus': 299,          # w.r.t. 2 VENUS BARYCENTER
42    'moon': 301,           # w.r.t. 3 EARTH BARYCENTER
43    'earth': 399,          # w.r.t. 3 EARTH BARYCENTER
44    'mars': 499,           # w.r.t. 4 MARS BARYCENTER
45    }
46
47
48class TestDAFBytesIO(TestCase):
49    def sample_daf(self):
50        word = Struct('d').pack
51        integer = Struct('i').pack
52        return BytesIO(b''.join([
53            # Record 1 - File Record
54            b'DAF/SPK ',
55            b'\x02\x00\x00\x00', # ND
56            b'\x03\x00\x00\x00', # NI
57            b'Internal Name'.ljust(60, b' '), # LOCIFN
58            b'\x03\x00\x00\x00', # FWARD
59            b'\x07\x00\x00\x00', # BWARD
60            b'\x01\x04\x00\x00', # FREE
61            b'LTL-IEEE', # LOCFMT
62            b'\0' * 603, # PRENUL
63            FTPSTR,
64            b'\0' * 297, # PSTNUL
65
66            # Record 2
67            b'Comment Record'.ljust(1024, b'\0'),
68
69            # Record 3 - first Summary Record
70            b''.join([
71                word(7), # next summary record
72                word(0), # previous summary record
73                word(1), # number of summaries
74                word(101),
75                word(202),
76                integer(303),
77                integer(1024 * 4 // 8 + 1), # Record 5 start
78                integer(1024 * 5 // 8),     # Record 5 end
79                integer(0),
80            ]).ljust(1024, b'\0'),
81
82            # Record 4 - first Name Record
83            b'Summary Name 1'.ljust(1024, b' '),
84
85            # Record 5
86            word(1001) * 128,
87
88            # Record 6
89            word(2002) * 128,
90
91            # Record 7 - second Summary Record
92            b''.join([
93                word(0), # next summary record
94                word(3), # previous summary record
95                word(1), # number of summaries
96                word(111),
97                word(222),
98                integer(333),
99                integer(1024 * 5 // 8 + 1), # Record 6 start
100                integer(1024 * 6 // 8),     # Record 6 end
101                integer(0),
102            ]).ljust(1024, b'\0'),
103
104            # Record 8 - second Name Record
105            b'Summary Name 2'.ljust(1024, b' '),
106        ]))
107
108    def test_header(self):
109        f = self.sample_daf()
110        d = DAF(f)
111        eq = self.assertEqual
112        eq(d.locidw, b'DAF/SPK')
113        eq(d.nd, 2)
114        eq(d.ni, 3)
115        eq(d.locifn_text, b'Internal Name')
116        eq(d.fward, 3)
117        eq(d.bward, 7)
118        eq(d.free, 0x401)
119        eq(d.locfmt, b'LTL-IEEE')
120
121    def test_segments(self):
122        f = self.sample_daf()
123        d = DAF(f)
124
125        summaries = list(d.summaries())
126        eq = self.assertEqual
127        eq(len(summaries), 2)
128        eq(summaries[0], (b'Summary Name 1', (101.0, 202.0, 303, 513, 640)))
129        eq(summaries[1], (b'Summary Name 2', (111.0, 222.0, 333, 641, 768)))
130
131        eq = self.assertSequenceEqual
132        eq(list(d.map(summaries[0][1])), [1001.0] * 128)
133        eq(list(d.map(summaries[1][1])), [2002.0] * 128)
134
135    def test_add_segment(self):
136        f = self.sample_daf()
137        d = DAF(f)
138
139        d.add_array(b'Summary Name 3', (121.0, 232.0, 343), [3003.0] * 128)
140
141        summaries = list(d.summaries())
142        eq = self.assertEqual
143        eq(len(summaries), 3)
144        eq(summaries[0], (b'Summary Name 1', (101.0, 202.0, 303, 513, 640)))
145        eq(summaries[1], (b'Summary Name 2', (111.0, 222.0, 333, 641, 768)))
146        eq(summaries[2], (b'Summary Name 3', (121.0, 232.0, 343, 1025, 1152)))
147
148        eq = self.assertSequenceEqual
149        eq(list(d.map(summaries[0][1])), [1001.0] * 128)
150        eq(list(d.map(summaries[1][1])), [2002.0] * 128)
151        eq(list(d.map(summaries[2][1])), [3003.0] * 128)
152
153    def test_add_segment_when_summary_block_is_full(self):
154        f = self.sample_daf()
155        d = DAF(f)
156
157        # Update n_summaries of final summary block to full.
158        d.file.seek(6 * 1024 + 16)
159        d.file.write(Struct('d').pack(d.summaries_per_record))
160
161        d.add_array(b'Summary Name 3', (121.0, 232.0, 343), [3003.0] * 200)
162
163        # Reset n_summaries of that block back to its real value.
164        d.file.seek(6 * 1024 + 16)
165        d.file.write(Struct('d').pack(1))
166
167        summaries = list(d.summaries())
168        eq = self.assertEqual
169        eq(len(summaries), 3)
170        eq(summaries[0], (b'Summary Name 1', (101.0, 202.0, 303, 513, 640)))
171        eq(summaries[1], (b'Summary Name 2', (111.0, 222.0, 333, 641, 768)))
172        eq(summaries[2], (b'Summary Name 3', (121.0, 232.0, 343, 1281, 1480)))
173
174        eq = self.assertSequenceEqual
175        eq(list(d.map(summaries[0][1])), [1001.0] * 128)
176        eq(list(d.map(summaries[1][1])), [2002.0] * 128)
177        eq(list(d.map(summaries[2][1])), [3003.0] * 200)
178
179
180class TestDAFRealFile(TestDAFBytesIO):
181    # Where "Real" = "written to disk with a real file descriptor
182    # instead of an in-memory BytesIO".
183
184    def sample_daf(self):
185        bytes_io = super(TestDAFRealFile, self).sample_daf()
186        f = tempfile.NamedTemporaryFile(mode='w+b', prefix='jplephem_test')
187        f.write(bytes_io.getvalue())
188        f.seek(0)
189        return f
190
191
192def fake_mmap_that_raises_OSError(*args, **kw):
193    raise OSError('mmap() not supported on this platform')
194
195class TestDAFRealFileWithoutMMap(TestDAFRealFile):
196    # Where "Real" = "written to disk with a real file descriptor
197    # instead of an in-memory BytesIO".  And we turn off mmap() to
198    # simulate platforms like pyodide.
199
200    def setUp(self):
201        self.mmap = mmap.mmap
202        mmap.mmap = fake_mmap_that_raises_OSError
203
204    def tearDown(self):
205        mmap.mmap = self.mmap
206
207
208class _CommonTests(object):
209
210    def check0(self, xyz, xyzdot=None):
211        eq = partial(self.assertAlmostEqual, delta=epsilon_m)
212        x, y, z = xyz
213        eq(x, 39705023.28)
214        eq(y, 131195345.65)
215        eq(z, 56898495.41)
216        if xyzdot is None:
217            return
218        dx, dy, dz = xyzdot
219        eq(dx, -2524248.19)
220        eq(dy, 619970.11)
221        eq(dz, 268928.26)
222
223    def check1(self, xyz, xyzdot=None):
224        eq = partial(self.assertAlmostEqual, delta=epsilon_m)
225        x, y, z = xyz
226        eq(x, -144692624.00)
227        eq(y, -32707965.14)
228        eq(z, -14207167.26)
229        if xyzdot is None:
230            return
231        dx, dy, dz = xyzdot
232        eq(dx, 587334.38)
233        eq(dy, -2297419.36)
234        eq(dz, -996628.74)
235
236    def test_scalar_tdb(self):
237        self.check0(self.position('earthmoon', 2414994.0))
238        self.check1(self.position('earthmoon', 2415112.5))
239
240    def test_scalar_tdb2(self):
241        self.check0(self.position('earthmoon', 2414990.0, 4.0))
242        self.check1(self.position('earthmoon', 2415110.0, 2.5))
243
244    def test_scalar_tdb_keyword(self):
245        self.check0(self.position('earthmoon', tdb=2414994.0))
246        self.check1(self.position('earthmoon', tdb=2415112.5))
247
248    def test_scalar_tdb2_keyword(self):
249        self.check0(self.position('earthmoon', tdb=2414990.0, tdb2=4.0))
250        self.check1(self.position('earthmoon', tdb=2415110.0, tdb2=2.5))
251
252    def check_2d_result(self, name, tdb, tdb2):
253        p = self.position(name, tdb + tdb2)
254        self.check0(p[:,0])
255        self.check1(p[:,1])
256
257        p = self.position(name, tdb, tdb2)
258        self.check0(p[:,0])
259        self.check1(p[:,1])
260
261        p, v = self.position_and_velocity(name, tdb + tdb2)
262        self.check0(p[:,0], v[:,0])
263        self.check1(p[:,1], v[:,1])
264
265        p, v = self.position_and_velocity(name, tdb, tdb2)
266        self.check0(p[:,0], v[:,0])
267        self.check1(p[:,1], v[:,1])
268
269    def test_array_tdb(self):
270        tdb = np.array([2414994.0, 2415112.5])
271        tdb2 = 0.0
272        self.check_2d_result('earthmoon', tdb, tdb2)
273
274    def test_array_tdb_scalar_tdb2(self):
275        tdb = np.array([2414991.5, 2415110.0])
276        tdb2 = 2.5
277        self.check_2d_result('earthmoon', tdb, tdb2)
278
279    def test_scalar_tdb_array_tdb2(self):
280        tdb = 2414990.0
281        d = 2415112.5 - tdb
282        tdb2 = np.array([4.0, d])
283        self.check_2d_result('earthmoon', tdb, tdb2)
284
285    def test_array_tdb_array_tdb2(self):
286        tdb = np.array([2414990.0, 2415110.0])
287        tdb2 = np.array([4.0, 2.5])
288        self.check_2d_result('earthmoon', tdb, tdb2)
289
290    def test_jitter(self):
291        usecond = 1e-6 / 24.0 / 3600.0
292        tdb = np.ones(12) * 2414998.0
293        tdb2 = np.linspace(1.0 * usecond, 2.0 * usecond, 12)
294        x, y, z = self.position('earthmoon', tdb, tdb2)
295        for component in x, y, z:
296            size = component[0]
297            relative_jitter = np.diff(np.diff(x)) / size
298            self.assertLess(max(abs(relative_jitter)), 3e-16)
299
300    def test_ephemeris_end_date(self):
301        x, y, z = self.position('earthmoon', self.jomega)
302        # These positions are actually from HORIZONS and thus DE431,
303        # hence the low precision match:
304        self.assertAlmostEqual(x, 1.442502234663646E+08, delta=1.0)
305        self.assertAlmostEqual(y, 3.690043031712407E+07, delta=1.0)
306        self.assertAlmostEqual(z, 1.599543968176661E+07, delta=1.0)
307
308    def test_too_early_date(self):
309        tdb = self.jalpha - 0.01
310        self.assertRaises(ValueError, self.position, 'earthmoon', tdb)
311
312    def test_too_late_date(self):
313        tdb = self.jomega + 16.01
314        self.assertRaises(ValueError, self.position, 'earthmoon', tdb)
315
316
317class SPKTests(_CommonTests, TestCase):
318
319    def setUp(self):
320        try:
321            self.spk = SPK.open('de421.bsp')
322        except IOError:
323            raise SkipTest('the "de421.bsp" SPK file is not available')
324        segment = self.spk[0,1]
325        self.jalpha = segment.start_jd
326        self.jomega = segment.end_jd
327
328    def tearDown(self):
329        self.spk.close()
330
331    def position(self, name, tdb, tdb2=0.0):
332        segment = self.spk[0, target_names[name]]
333        return segment.compute(tdb, tdb2)
334
335    def position_and_velocity(self, name, tdb, tdb2=0.0):
336        segment = self.spk[0, target_names[name]]
337        return segment.compute_and_differentiate(tdb, tdb2)
338
339    def test_segment_with_only_two_coefficients(self):
340        tdb = 2414990.0
341        tup = target_names['mercury barycenter'], target_names['mercury']
342        segment = self.spk[tup]
343        segment.compute_and_differentiate(tdb)
344
345    def test_str(self):
346        str(self.spk)  # just to confirm it does not raise an exception
347        segment = self.spk[0,4]
348        self.assertEqual(str(segment), segment.describe(verbose=False))
349        self.assertEqual(segment.describe(verbose=False),
350  '2414864.50..2471184.50  Type 2  Solar System Barycenter (0) -> Mars Barycenter (4)')
351        self.assertEqual(segment.describe(verbose=True),
352  '2414864.50..2471184.50  Type 2  Solar System Barycenter (0) -> Mars Barycenter (4)'
353  '\n  frame=1 source=DE-0421LE-0421')
354
355    def test_loading_array(self):
356        segment = self.spk[0,4]
357        initial_epoch, interval_length, coefficients = segment.load_array()
358        self.assertEqual(coefficients.shape, (3, 1760, 11))
359
360    def test_out_of_range_dates(self):
361        segment = self.spk[0,4]
362        tdb = np.array([-1e3, 0, +1e5]) + 2414990.0
363        try:
364            segment.compute_and_differentiate(tdb)
365        except OutOfRangeError as e:
366            self.assertEqual(str(e), 'segment only covers dates'
367                             ' 2414864.5 through 2471184.5')
368            self.assertIs(type(e.out_of_range_times), np.ndarray)
369            self.assertEqual(list(e.out_of_range_times), [True, False, True])
370
371class LegacyTests(_CommonTests, TestCase):
372
373    def setUp(self):
374        try:
375            import de421
376        except ImportError:
377            raise SkipTest('the "de421" ephemeris package has not been'
378                           ' installed with "pip install de421"')
379        self.eph = Ephemeris(de421)
380        self.jalpha = self.eph.jalpha
381        self.jomega = self.eph.jomega
382
383    def position(self, name, tdb, tdb2=0.0):
384        return self.eph.position(name, tdb, tdb2)
385
386    def position_and_velocity(self, name, tdb, tdb2=0.0):
387        return self.eph.position_and_velocity(name, tdb, tdb2)
388
389    def test_names(self):
390        self.assertEqual(self.eph.names,  (
391            'earthmoon', 'jupiter', 'librations', 'mars', 'mercury',
392            'moon', 'neptune', 'nutations', 'pluto', 'saturn', 'sun',
393            'uranus', 'venus',
394            ))
395
396    def test_legacy_compute_method(self):
397        pv = self.eph.compute('earthmoon', 2414994.0)
398        self.check0(pv[:3], pv[3:])
399        pv = self.eph.compute('earthmoon', np.array([2414994.0, 2415112.5]))
400        self.check0(pv[:3,0], pv[3:,0])
401        self.check1(pv[:3,1], pv[3:,1])
402
403    def test_ephemeris_end_date(self):
404        x, y, z = self.position('earthmoon', self.jomega)
405        self.assertAlmostEqual(x, -94189805.73967789, delta=epsilon_m)
406        self.assertAlmostEqual(y, 1.05103857e+08, delta=1.0)
407        self.assertAlmostEqual(z, 45550861.44383482, delta=epsilon_m)
408
409
410class NAIF_DAF_Tests(TestCase):
411
412    def test_single_position(self):
413        with SPK(NAIF_DAF(open('de405.bsp', 'rb'))) as kernel:
414            x, y, z = kernel[0,4].compute(2457061.5)
415            # Expect rough agreement with a DE430 position from our README:
416            self.assertAlmostEqual(x, 2.05700211e+08, delta=2.0)
417            self.assertAlmostEqual(y, 4.25141646e+07, delta=2.0)
418            self.assertAlmostEqual(z, 1.39379183e+07, delta=2.0)
419
420
421class CommandLineTests(TestCase):
422    maxDiff = 9999
423
424    def test_comment_command(self):
425        output = commandline.main(['comment', 'de405.bsp'])
426        self.assertEqual(output[:30], '; de405.bsp LOG FILE\n;\n; Creat')
427        self.assertEqual(output[-30:], "rom Standish's DE405 memo <<<\n")
428
429    def test_daf_command(self):
430        self.assertEqual(commandline.main(['daf', 'de405.bsp']), """\
431 1 DE-405 -1577879958.8160586 1577880064.1839132 1 0 1 2 1409 202316
432 2 DE-405 -1577879958.8160586 1577880064.1839132 2 0 1 2 202317 275376
433 3 DE-405 -1577879958.8160586 1577880064.1839132 3 0 1 2 275377 368983
434 4 DE-405 -1577879958.8160586 1577880064.1839132 4 0 1 2 368984 408957
435 5 DE-405 -1577879958.8160586 1577880064.1839132 5 0 1 2 408958 438653
436 6 DE-405 -1577879958.8160586 1577880064.1839132 6 0 1 2 438654 464923
437 7 DE-405 -1577879958.8160586 1577880064.1839132 7 0 1 2 464924 487767
438 8 DE-405 -1577879958.8160586 1577880064.1839132 8 0 1 2 487768 510611
439 9 DE-405 -1577879958.8160586 1577880064.1839132 9 0 1 2 510612 533455
44010 DE-405 -1577879958.8160586 1577880064.1839132 10 0 1 2 533456 613364
44111 DE-405 -1577879958.8160586 1577880064.1839132 301 3 1 2 613365 987780
44212 DE-405 -1577879958.8160586 1577880064.1839132 399 3 1 2 987781 1362196
44313 DE-405 -1577879958.8160586 1577880064.1839132 199 1 1 2 1362197 1362208
44414 DE-405 -1577879958.8160586 1577880064.1839132 299 2 1 2 1362209 1362220
44515 DE-405 -1577879958.8160586 1577880064.1839132 499 4 1 2 1362221 1362232
446""")
447
448    def test_spk_command(self):
449        self.assertEqual(commandline.main(['spk', 'de405.bsp']), """\
450File type NAIF/DAF and format BIG-IEEE with 15 segments:
4512433282.50..2469807.50  Type 2  Solar System Barycenter (0) -> Mercury Barycenter (1)
4522433282.50..2469807.50  Type 2  Solar System Barycenter (0) -> Venus Barycenter (2)
4532433282.50..2469807.50  Type 2  Solar System Barycenter (0) -> Earth Barycenter (3)
4542433282.50..2469807.50  Type 2  Solar System Barycenter (0) -> Mars Barycenter (4)
4552433282.50..2469807.50  Type 2  Solar System Barycenter (0) -> Jupiter Barycenter (5)
4562433282.50..2469807.50  Type 2  Solar System Barycenter (0) -> Saturn Barycenter (6)
4572433282.50..2469807.50  Type 2  Solar System Barycenter (0) -> Uranus Barycenter (7)
4582433282.50..2469807.50  Type 2  Solar System Barycenter (0) -> Neptune Barycenter (8)
4592433282.50..2469807.50  Type 2  Solar System Barycenter (0) -> Pluto Barycenter (9)
4602433282.50..2469807.50  Type 2  Solar System Barycenter (0) -> Sun (10)
4612433282.50..2469807.50  Type 2  Earth Barycenter (3) -> Moon (301)
4622433282.50..2469807.50  Type 2  Earth Barycenter (3) -> Earth (399)
4632433282.50..2469807.50  Type 2  Mercury Barycenter (1) -> Mercury (199)
4642433282.50..2469807.50  Type 2  Venus Barycenter (2) -> Venus (299)
4652433282.50..2469807.50  Type 2  Mars Barycenter (4) -> Mars (499)
466""")
467
468
469def load_tests(loader, tests, ignore):
470    """Run our main documentation as a test."""
471
472    # If we are running in CI, where we test against an old version of
473    # NumPy, skip the doctests since NumPy will print whitespace
474    # differently (and worse).
475    version = tuple(int(s) for s in np.__version__.split('.'))
476    if version < (1, 17):
477        return tests
478
479    # Python 2.6 formats floating-point numbers a bit differently and
480    # breaks the doctest.
481    if sys.version_info <= (2, 6):
482        return tests
483
484    tests.addTests(DocTestSuite('jplephem', optionflags=ELLIPSIS))
485    return tests
486