fix snapshot outputting for mockstream

gala/dynamics/mockstream/mockstream.pyx

@@ -1,4 +1,4 @@
-# cython: boundscheck=False
+ # cython: boundscheck=False
 # cython: debug=False
 # cython: nonecheck=False
 # cython: cdivision=True
@@ -204,10 +204,10 @@ cpdef mockstream_dop853_animate(nbody, double[::1] t,
         CPotential *c_particle_potentials[MAX_NBODY]
 
         # Snapshotting:
-        int noutput_times = ntimes // output_every
+        int noutput_times = (ntimes-1) // output_every + 1
         double[::1] output_times
 
-    if ntimes % output_every != 0:
+    if (ntimes-1) % output_every != 0:
         noutput_times += 1 # +1 for final conditions
 
     output_times = np.zeros(noutput_times)
@@ -278,23 +278,26 @@ cpdef mockstream_dop853_animate(nbody, double[::1] t,
 
     j = 1 # output time index
     for i in range(1, ntimes):
+        # print(i, j, n,
+        #       len(t), len(nstream), len(output_times))
+
         dop853_step(&cp, &cf, <FcnEqDiff> Fwrapper_direct_nbody,
                     &w[0, 0], t[i-1], t[i], dt0,
                     ndim, nbodies+n, nbodies, args,
                     atol, rtol, nmax)
 
         PyErr_CheckSignals()
 
-        if (i % output_every) == 0:
+        n += nstream[i]
+
+        if (i % output_every) == 0 or i == ntimes-1:
             output_times[j] = t[i]
             stream_g['pos'][:, j, :n] = np.array(w[nbodies:nbodies+n, :]).T[:3]
             stream_g['vel'][:, j, :n] = np.array(w[nbodies:nbodies+n, :]).T[3:]
             nbody_g['pos'][:, j, :n] = np.array(w[:nbodies, :]).T[:3]
             nbody_g['vel'][:, j, :n] = np.array(w[:nbodies, :]).T[3:]
             j += 1
 
-        n += nstream[i]
-
     for g in [stream_g, nbody_g]:
         d = g.create_dataset('time', data=np.array(output_times))
         d.attrs['unit'] = str(nbody.units['time'])

gala/dynamics/mockstream/mockstream_generator.py

@@ -11,6 +11,7 @@
 
 __all__ = ['MockStreamGenerator']
 
+
 class MockStreamGenerator:
 
     def __init__(self, df, hamiltonian, progenitor_potential=None):
@@ -154,9 +155,7 @@ def run(self, prog_w0, prog_mass, nbody=None,
             bursts of particles at certain times (e.g., pericenter).
         output_every : int (optional)
             Controls whether to output snapshots of the stream particle orbits.
-            This is incorporated after ``release_every``, so if
-            ``release_every=4`` and ``output_every=4``, the snapshots will be
-            saved every 16 steps of the global time array, for example.
+            This is relative to the global time array.
         output_filename : str (optional)
             The path to the HDF5 file to be generated by the snapshotting.
         check_filesize : bool (optional)
@@ -190,7 +189,7 @@ def run(self, prog_w0, prog_mass, nbody=None,
         if t[1] < t[0]:
             nbody_orbits = nbody_orbits[::-1]
 
-            # TODO: this could be cleaed up...
+            # TODO: this could be cleaned up...
             nbody0 = DirectNBody(
                 nbody_orbits[0], prog_nbody.particle_potentials,
                 external_potential=self.hamiltonian.potential,
@@ -199,7 +198,7 @@ def run(self, prog_w0, prog_mass, nbody=None,
         else:
             nbody0 = prog_nbody
 
-        prog_orbit = nbody_orbits[:, 0] # Note: Progenitor must be idx 0!
+        prog_orbit = nbody_orbits[:, 0]  # Note: Progenitor must be idx 0!
         orbit_t = prog_orbit.t.decompose(units).value
 
         # Generate initial conditions from the DF
@@ -218,7 +217,7 @@ def run(self, prog_w0, prog_mass, nbody=None,
         for t1, n in zip(unq_t1s, nstream):
             all_nstream[orbit_t == t1] = n
 
-        if output_every is None: # store snapshots
+        if output_every is None:  # store snapshots
             raw_nbody, raw_stream = mockstream_dop853(
                 nbody0, orbit_t[all_nstream != 0], w0, unq_t1s,
                 all_nstream[all_nstream != 0].astype('i4'))
@@ -228,7 +227,7 @@ def run(self, prog_w0, prog_mass, nbody=None,
                                  "pass in a filename to store the snapshots in")
 
             raw_nbody, raw_stream = mockstream_dop853_animate(
-                nbody0, orbit_t, w0, nstream.astype('i4'),
+                nbody0, orbit_t, w0, all_nstream.astype('i4'),
                 output_every=output_every, output_filename=output_filename,
                 check_filesize=check_filesize, overwrite=overwrite)
 

gala/dynamics/mockstream/tests/test_mockstream.py

@@ -1,4 +1,5 @@
 import os
+import itertools
 
 # Third-party
 import astropy.units as u
@@ -95,9 +96,10 @@ def test_run():
     # TODO: add nbody test
 
 
-# @pytest.mark.skipif('CI' in os.environ,
-#                     reason="For some reason, doesn't work on Travis/CI")
-def test_animate(tmpdir):
+@pytest.mark.parametrize(
+    'dt, output_every, release_every, n_particles, trail',
+    list(itertools.product([1, -1], [1, 2], [1, 4], [1, 4], [True, False])))
+def test_animate(tmpdir, dt, output_every, release_every, n_particles, trail):
     import h5py
 
     potential = NFWPotential.from_circular_velocity(v_c=0.2, r_s=20.,
@@ -108,21 +110,30 @@ def test_animate(tmpdir):
     mass = 2.5e4 * u.Msun
 
     # The basic run:
-    df = FardalStreamDF(trail=False)
+    df = FardalStreamDF(trail=trail)
     gen = MockStreamGenerator(df=df, hamiltonian=H)
 
     filename = os.path.join(str(tmpdir), "test.hdf5")
     nsteps = 16
-    output_every = 2
-    stream, _ = gen.run(w0, mass, dt=-1., n_steps=nsteps,
+    stream, _ = gen.run(w0, mass, dt=dt, n_steps=nsteps,
+                        release_every=release_every,
+                        n_particles=n_particles,
                         output_every=output_every, output_filename=filename,
                         overwrite=True)
 
     with h5py.File(filename, mode='r') as f:
         stream_orbits = Orbit.from_hdf5(f['stream'])
         nbody_orbits = Orbit.from_hdf5(f['nbody'])
 
-    assert stream_orbits.shape == (1 + nsteps // output_every, nsteps+1)
+    noutput_times = 1 + nsteps // output_every
+    if nsteps % output_every != 0:
+        noutput_times += 1
+
+    tail_n_particles = (1 + int(trail)) * n_particles
+    expected_shape = (noutput_times,
+                      tail_n_particles * (nsteps // release_every + 1))
+
+    assert stream_orbits.shape == expected_shape
     assert np.isfinite(stream_orbits[:, 0].xyz).all()
     assert np.isfinite(stream_orbits[:, 0].v_xyz).all()