Don’t pre-build vector functions for eph segments

To reduce the expense of looking up a planet in an ephemeris, each
ephemeris was doing a single scan of its `.segments` in `__init__()` and
caching the resulting vector functions.  But this meant that users could
not usefully edit the `.segments` list — for example, to combine all the
segments from two ephemeris files.  So let’s do everything on-the-fly.

Author: brandon-rhodes

skyfield/jpllib.py

@@ -73,19 +73,12 @@ def __init__(self, path):
         self.segments = [SPICESegment(self, s) for s in self.spk.segments]
         self.comments = self.spk.comments  # deprecated pass-through method
 
-        # Pre-compute which segments lead to which targets.
-        d = defaultdict(list)
-        for segment in self.segments:
-            d[segment.target].append(segment)
-
-        # Go ahead and build a vector function for each target.
-        self._vector_functions = {
-            target: segments[0] if len(segments) == 1 else Stack(segments)
-            for target, segments in d.items()
-        }
-
     def __repr__(self):
-        return '<{0} {1!r}>'.format(type(self).__name__, self.path)
+        paths = sorted({s.ephemeris.path for s in self.segments})
+        return '<{} {}>'.format(
+            type(self).__name__,
+            ' '.join(repr(path) for path in paths),
+        )
 
     def __str__(self):
         lines = []
@@ -94,8 +87,9 @@ def __str__(self):
         ephemeris = start = end = None
         for s in self.segments:
             if ephemeris is not s.ephemeris:
+                word = 'Segments' if (ephemeris is None) else 'And'
                 ephemeris = s.ephemeris
-                a('Segments from kernel file {!r}:'.format(ephemeris.filename))
+                a('{} from kernel file {!r}:'.format(word, ephemeris.filename))
             spk = s.spk_segment
             if start != spk.start_jd or end != spk.end_jd:
                 start, end = spk.start_jd, spk.end_jd
@@ -170,17 +164,28 @@ def decode(self, name):
 
     def __getitem__(self, target):
         """Return a vector function for computing the location of `target`."""
+        target_arg = target
         target = self.decode(target)
-        vector_functions = self._vector_functions
-        vf = vector_functions[target]
-        if vf.center == 0:
-            return vf
-        vfs = [vf]
-        center = vf.center
-        while center in vector_functions:
-            vf = vector_functions[center]
+        center = target
+        vfs = []
+
+        while center != 0:
+            matches = [s for s in self.segments if s.target == center]
+            if not matches:
+                raise KeyError(
+                    'the segments of this ephemeris cannot connect the Solar'
+                    ' System Barycenter to the target {!r}'.format(target_arg)
+                )
+            elif len(matches) == 1:
+                vf = matches[0]
+            else:
+                vf = Stack(matches)
             vfs.append(vf)
             center = vf.center
+
+        if len(vfs) == 1:
+            return vfs[0]
+
         return VectorSum(center, target, tuple(reversed(vfs)))
 
     def __contains__(self, name_or_code):

skyfield/tests/test_jpllib.py

@@ -32,7 +32,7 @@ def test_multiple_non_overlapping_segments_per_target():
 
 # Verify that ephemeris objects let their segments be edited.
 
-def test_removing_segments_from_jpl_ephemeris():
+def test_removing_segments_from_ephemeris():
     eph = load('de421.bsp')
     eph.segments = [s for s in eph.segments if s.target in (3, 4)]
 
@@ -55,5 +55,50 @@ def test_removing_segments_from_jpl_ephemeris():
       0 -> 4    SOLAR SYSTEM BARYCENTER -> MARS BARYCENTER"""
 
     assert type(eph[4]).__name__ == 'ChebyshevPosition'
-    with assert_raises(KeyError):
+    with assert_raises(KeyError, 'is missing 5'):
         eph[5]
+
+def test_adding_segments_to_ephemeris():
+    eph = load('de405.bsp')
+    eph.segments = [s for s in eph.segments if s.target == 3]
+
+    eph2 = load('de421.bsp')
+    eph.segments.extend((s for s in eph2.segments if s.target in (301, 399)))
+
+    assert len(eph.segments) == 3
+    assert 2 not in eph
+    assert 3 in eph
+    assert 301 in eph
+    assert eph.codes == {0, 3, 301, 399}
+    assert eph.names() == {
+        0: ['SOLAR_SYSTEM_BARYCENTER', 'SSB', 'SOLAR SYSTEM BARYCENTER'],
+        3: ['EARTH_BARYCENTER', 'EMB', 'EARTH MOON BARYCENTER',
+            'EARTH-MOON BARYCENTER', 'EARTH BARYCENTER'],
+        301: ['MOON'], 399: ['EARTH'],
+    }
+
+    assert repr(eph) == "<SpiceKernel 'de405.bsp' 'de421.bsp'>"
+    assert str(eph) == """\
+Segments from kernel file 'de405.bsp':
+  JD 2305424.50 - JD 2525008.50  (1599-12-08 through 2201-02-19)
+      0 -> 3    SOLAR SYSTEM BARYCENTER -> EARTH BARYCENTER
+And from kernel file 'de421.bsp':
+  JD 2414864.50 - JD 2471184.50  (1899-07-28 through 2053-10-08)
+      3 -> 301  EARTH BARYCENTER -> MOON
+      3 -> 399  EARTH BARYCENTER -> EARTH"""
+
+    vs = eph[399]
+    assert type(vs).__name__ == 'VectorSum'
+    assert str(vs) == """\
+Sum of 2 vectors:
+ 'de405.bsp' segment 0 SOLAR SYSTEM BARYCENTER -> 3 EARTH BARYCENTER
+ 'de421.bsp' segment 3 EARTH BARYCENTER -> 399 EARTH"""
+
+    with assert_raises(KeyError):
+        eph[4]
+
+def test_ephemeris_lacking_segments_to_connect_to_barycenter():
+    eph = load('de421.bsp')
+    eph.segments = [s for s in eph.segments if s.target == 399]
+    with assert_raises(KeyError, "Barycenter to the target 'Earth'"):
+        eph['Earth']