Smooth stream track for streamspraydf models

[jobovy]

Summary

• Adds StreamTrack (and StreamTrackPair for tail='both') — a lightweight smooth phase-space track parameterized by the progenitor's time coordinate tp in [-tdisrupt, 0]. Exposes Orbit-like accessors for Cartesian, cylindrical, and heliocentric coordinates (x/y/z/vx/vy/vz, R/vR/vT/phi, ra/dec/dist/ll/bb/pmra/pmdec/pmll/pmbb/vlos) plus a cov(tp) and plot(d1, d2, spread=...).
• Adds basestreamspraydf.streamTrack(n=5000, tail=..., smoothing=..., niter=0, particles=..., order=2) which samples particles (or reuses pre-computed ones) and fits the track. The track is the smoothed binned mean of particles as a function of tp=-dt with a smoothed 6×6 covariance.
• Tests, API reference docs, and a new notebook section in the spraydf tutorial.

A comment in streamdf.py points at StreamTrack as a candidate shared backend for a future follow-up PR that could simplify streamdf's track plumbing.

Test plan

• 11 new unit tests in tests/test_streamspraydf.py pass (progenitor recovery, sample consistency, interface, PSD covariance, both arms, iteration, chen24+fardal15, center/satellite streams, physical units, particle reuse, plot smoke).
• Existing test_streamspraydf.py tests still pass (28 total passing locally).
• Notebook section uses existing spdf_both fixture; cells run interactively against a local execution.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 99.91%. Comparing base (9494856) to head (a6712e8).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff            @@
##             main     #861    +/-   ##
========================================
  Coverage   99.91%   99.91%            
========================================
  Files         224      225     +1     
  Lines       32912    33619   +707     
  Branches      710      710            
========================================
+ Hits        32885    33592   +707     
  Misses         27       27            

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jobovy]

StreamTrack implementation alternatives — comparison summary

Seven alternative branches were explored, each changing one axis of the algorithm. All branch off streamspray-track and are pushed (no PRs). Comparison plots (Bovy14 + Pal 5 + warm 10^7 Msun streams, aligned phi_1/phi_2 coordinates) are in the repo root as alt_*.png files.

Main branch (streamspray-track)

Closest-point projection (3D position) onto a finely-integrated, short-range progenitor orbit (0.03*tdisrupt). Bin-then-smooth offsets-from-progenitor with UnivariateSpline (chi-square s=N). Percentile-trimmed tp_grid.

Alternatives tested

Branch	What it changes	Cold stream (Bovy14)	Pal 5
alt-gcv	GCV auto-smoothing (make_smoothing_spline)	Eliminates leading-arm wiggle	Clean
alt-no-binning	Per-particle splines, no binning	Slight tip oscillation	Small spikes at tips
alt-rotating-frame	Smooth in progenitor-aligned rotating frame	Clean	Clean
alt-6d-closest	6D (position+velocity) closest-point	Clean	Clean
alt-strip-time-affine	tp = arm_sign × dt (no closest-point)	Dramatic wrap-aliasing failure	Same
alt-auto-timerange	Data-driven track_time_range from particle spread	Clean	Clean
alt-kdtree	cKDTree for closest-point (speed)	Identical to main	Identical

Key findings

1. alt-strip-time-affine (tp=dt directly) fails catastrophically for multi-wrap streams. The track loops across the sky multiple times. This validates the closest-point matching approach.

2. alt-gcv (GCV smoothing) fixes the leading-arm wiggle visible in the main branch's Bovy14 fardal15 track. The main branch's chi-square s=N criterion slightly under-smooths; GCV auto-selects a better lambda.

3. alt-auto-timerange is essential for warm streams (see next comment).

4. alt-no-binning, alt-rotating-frame, alt-6d-closest, alt-kdtree produce tracks within ≤0.05 galpy units of main for cold streams — the algorithmic choices they test are secondary for typical use.

Recommendation

Combine: auto-timerange (scales with stream width) + GCV smoothing (eliminates under-smoothing) + 6D matching (robustness for eccentric orbits) + boundary pile-up exclusion + percentile trimming. Require scipy >= 1.10 (released Jan 2023).

Comparisons

Main branch (streamspray-track, PR #861)

Closest-point projection of each particle onto a finely-integrated, short-range (0.03*tdisrupt by default) progenitor orbit that spans both sides of tp=0. Bin-then-smooth offsets-from-progenitor with UnivariateSpline (chi-square-like auto-s). Covariance smoothed per-entry and PSD-projected.

The public tp_grid is trimmed to the 99th percentile of the
assigned tp values (symmetric 1st percentile on the trailing side),
so a handful of outlier particles matched to sparsely-populated
boundary bins can't drag the smoothing spline around at the tips.
trim_percentile is a local constant in StreamTrack.__init__ — we can
expose it as a kwarg later if there's demand.

Bovy14 stream (LogarithmicHaloPotential, q=0.9, tdisrupt=4.5 Gyr):

Pal 5 stream (MWPotential2014, Orbit.from_name("Pal 5"), tdisrupt=5 Gyr):

All 7 alternatives below inherit the same trim and can be toggled by git checkout alt-<name>. Each plot now has two panels: an aligned (phi_1, phi_2) view with the alt's track on top of the sample cloud and the streamspray-track reference as dashed black lines, plus a difference curve per arm against the saved reference.

What the bottom panel plots. The y-axis on the difference panel is the 3D Euclidean distance between the two tracks' galactocentric Cartesian positions at each tp:

d(tp) = ||(x, y, z)_alt(tp) − (x, y, z)_main(tp)||

in galpy internal length units (multiply by ro ≈ 8 for kpc). Not phi_2, not a single coordinate. Velocities are ignored in the diff metric; they're smoothed the same way as position, so positional agreement is a good proxy.

1. alt-gcv — auto-smoothing via GCV

scipy.interpolate.make_smoothing_spline replaces UnivariateSpline; the smoothing parameter is selected automatically by Generalized Cross Validation, so the user doesn't tune s.

• Pros: no smoothness knob; data-driven bias/variance tradeoff.
• Cons: requires scipy>=1.10; slower per fit.

2. alt-no-binning — per-particle smoothing splines

Skip binning entirely. Each particle contributes one (tp_i, offset_i) point to six cubic smoothing splines; covariance comes from smoothing splines on residual outer-products.

• Pros: no ntp knob, no empty-bin edge cases, adapts to uneven particle density along tp.
• Cons: fits N points rather than ntp (slower for large n); FITPACK occasionally warns when s is mis-estimated from pairwise noise.

3. alt-rotating-frame — smooth in the progenitor-aligned frame

At each particle's tp, rotate the raw offset into a progenitor-aligned frame (L along z, progenitor at +x). Smooth in that frame, then rotate back. This was the plan's original proposal.

• Pros: rotated offsets carry clean physical meaning (along-stream / transverse / vertical); dynamic range of the smoothed signal is minimal.
• Cons: per-tp rotation matrices add bookkeeping; interpolated matrices need SVD re-orthogonalization.

4. alt-6d-closest — closest-point in 6D instead of 3D

Match particles to the progenitor orbit in full 6D phase space (position + velocity, both in galpy internal units) rather than 3D position alone.

• Pros: unambiguous at orbit self-intersections (same xyz, different v); cleaner tp assignments near pericenter/apocenter.
• Cons: larger distance-matrix memory; the position/velocity weighting is fixed by the choice of units rather than physically motivated.

5. alt-strip-time-affine — tp = arm_sign * dt

Simplest possible mapping from particles to the affine parameter: use the known stripping time dt directly (positive for leading, negative for trailing). track_time_range defaults to tdisrupt so the progenitor is integrated over the full [-tdisrupt, +tdisrupt].

• Pros: fully deterministic, no closest-point failure modes, exact per-particle affine value.
• Cons: tp and along-stream position are correlated with scatter (ΔΩ varies per particle), so the mean-track at a given tp isn't exactly on the stream's median curve. Three tests are relaxed / one skipped on this branch because their tp semantics assume the closest-point mapping.

6. alt-auto-timerange — adaptive track_time_range from particle spread

Replace the fixed 0.03*tdisrupt default with a data-driven estimate: measure the stream's spatial extent from a probe sample, convert to orbital-time via progenitor speed, pad by 4x, clamp to [1, tdisrupt].

• Pros: adapts to narrow vs wide streams automatically; safe across a wider range of progenitor/stream regimes.
• Cons: needs one small extra sample draw when particles= not supplied; speed metric can be fragile for highly eccentric orbits.

7. alt-kdtree — KD-tree closest-point matching

Replace the O(N*M) pairwise distance matrix with a scipy.spatial.cKDTree query. When dt/arm-sign masks restrict the allowed neighbors per particle, query K nearest and pick the closest allowed one, growing K 4x per pass until every point has a match.

• Pros: sublinear per-query complexity; no O(N*M) memory footprint (matters for very large N or denser progenitor grids).
• Cons: a Python-level fallback loop runs when masks are restrictive (still fast for N up to ~1e5).

[jobovy]

Warm-stream comparison (10^7 Msun progenitor)

To test robustness for wider streams (dwarf-galaxy-mass progenitors), we ran all alternatives with the Bovy14 orbit but 1000× larger progenitor mass (10^7 Msun instead of 10^4 Msun). This produces a much fatter stream with ~10× larger tidal radius and velocity kicks.

Results

The main branch (3D closest-point, track_time_range = 0.03*tdisrupt = 3.8) fails: the leading arm track flies off to phi_2 ~ -35, far from the sample cloud. The root cause is insufficient progenitor orbit coverage: with only 3.8 galpy time units (~0.3 orbital periods), far-flung warm-stream particles can't find a close match on the short progenitor arc and pile up at the boundary.

Branch	Warm stream result	Why
main (3D, short range)	Broken — leading arm flies off	Progenitor arc too short for wide stream
alt-6d-closest (6D, short range)	Still problematic	Same short arc; 6D can't help if orbit isn't long enough
alt-gcv (GCV, short range)	Same failure	Better smoothing can't fix bad tp assignment
alt-rotating-frame (short range)	Same failure	Better frame can't fix bad tp assignment
alt-auto-timerange (3D, data-driven range)	Works — track follows cloud	Auto range = 13.6 (~1.1 orbital periods), covers all particles

Key numbers

• Cold stream: d_max = 1.7, auto range = 5.2 (1.4× main)
• Warm stream: d_max = 4.4, auto range = 13.6 (3.6× main)

The auto-timerange scales naturally with stream width because it measures the actual particle spread. The fixed 0.03*tdisrupt doesn't adapt.

Conclusion

For warm streams, progenitor orbit coverage is the primary bottleneck, not matching dimension or smoothing method. The auto-timerange fix is essential. 6D matching and GCV smoothing are secondary improvements that help with cold-stream edge cases (orbit self-intersection, under-smoothing) but cannot compensate for insufficient orbit coverage.

Boundary pile-up

Even with auto-timerange, particles at the edge of the progenitor arc are mismatched (their true closest point is outside the arc). These create a "pile-up" at the extreme tp values that biases the track. Fix: exclude particles whose tp_assign equals the far boundary of track_t_grid before fitting.

Example warm stream

[jobovy]

Warm-stream comparison (10^7 Msun progenitor, TriaxialNFW not included here)

Bovy14 orbit in LogarithmicHaloPotential(q=0.9), tdisrupt=4.5 Gyr, fardal15spraydf. Six projections: aligned sky (phi_1, phi_2), galactocentric (x, y), cylindrical (R, z), equatorial (RA, Dec), Galactic (l, v_los), heliocentric (dist, pm_ra).

Cold stream (10^4 Msun) — main branch

Warm stream (10^7 Msun) — main branch

Key findings

• Cold stream: tracks follow the sample cloud cleanly in all six projections. GCV smoothing produces a smooth curve with no leading-arm wiggle.
• Warm stream: tracks follow the sample cloud in all projections. The auto-timerange correctly picks track_time_range ≈ 27 (vs 3.8 for cold), covering the larger spatial extent. At every tp evaluation point, 74–436 sample particles are within 3 kpc of the track — confirming the track is correctly tracing through the thick cloud.
• Some sky-coordinate panels show straight-line artifacts where matplotlib connects tp-consecutive points across the RA=0/360 wrap boundary. This is a visualization issue, not a track error.

Diagnostics

	Cold (10^4)	Warm (10^7)
auto track_time_range	~3.4	~27
leading tp range	[0, 1.1]	[0, 11.8]
particles after boundary exclusion	3000/3000	3000/3000

[jobovy]

Triaxial NFW comparison: main vs stream-orbit v1 vs v2

TriaxialNFWPotential(normalize=1, a=20/8, b=0.8, c=0.6), Bovy14 orbit, tdisrupt=2 Gyr, n=1000 particles. Three projections: aligned sky, (x, y), (R, z).

Cold stream (10^4 Msun)

Main (progenitor orbit only):

v1 (one-shot orbit from raw particle-mean tip IC):

v2 (two-step orbit from smoothed track tip IC):

Cold verdict: All three nearly identical. The cold stream stays close to the progenitor orbit even in the triaxial potential.

---

Warm stream (10^7 Msun)

Main (progenitor orbit only):

v1 (one-shot orbit from raw particle-mean tip IC):

v2 (two-step orbit from smoothed track tip IC):

Warm verdict:

• Main: both arms follow the thick cloud in all projections. Clean.
• v1: trailing OK, but leading arm shoots off to R~26 in (R,z) — the noisy particle-mean tip IC leads to an orbit that diverges from the actual stream.
• v2: matches main quality. The smoothed-track tip IC avoids the v1 divergence.

---

Summary

	Cold (1e4)	Warm (1e7)
Main (progenitor only)	Clean	Clean
v1 (particle-mean tip)	Clean	Leading diverges in (R,z)
v2 (smoothed-track tip)	Clean	Clean, matches main

The main branch already handles the triaxial warm case well — the progenitor orbit is a good enough reference. v2 matches main; v1 still fails from IC noise. The stream-orbit refinement might become more impactful for more extreme triaxiality (b < 0.6) or very long tdisrupt where cumulative orbital precession builds up.

[jobovy]

Triaxial NFW comparison (updated: 5 Gyr tdisrupt)

TriaxialNFWPotential(normalize=1, a=20/8, b=0.8, c=0.6), Bovy14 orbit, tdisrupt=5 Gyr, n=1000 particles.

Cold stream (10^4 Msun)

Main (progenitor orbit only):

v1 (one-shot orbit from raw particle-mean tip IC):

v2 (two-step orbit from smoothed track tip IC):

---

Warm stream (10^7 Msun)

Main (progenitor orbit only):

v1 (one-shot orbit from raw particle-mean tip IC):

v2 (two-step orbit from smoothed track tip IC):

---

Summary

	Cold (1e4)	Warm (1e7)
Main (progenitor only)	Clean	Clean
v1 (particle-mean tip)	Clean	Leading arm diverges
v2 (smoothed-track tip)	Clean	Clean, matches main

With 5 Gyr of disruption in the triaxial potential, the main branch still handles both cold and warm streams well. v1's noisy IC diverges for the warm case; v2's smoothed IC fixes this and matches main quality. The stream-orbit refinement doesn't visibly improve over using the progenitor orbit directly at this level of triaxiality (b=0.8, c=0.6).

[jobovy]

Status update: warm-stream testing, combined improvements, and stream-orbit exploration

Warm-stream testing (10^7 Msun progenitor)

Stress-tested the algorithm with a 1000x larger progenitor mass. The fixed track_time_range = 0.03 * tdisrupt was too short for warm streams — particles spread far from the progenitor couldn't find matches on the short progenitor arc, causing boundary pile-up and a broken track. Fix: auto-timerange from particle spread (8 * d_max / |v_prog|), which scales naturally with stream width.

Combined implementation

Merged four orthogonal improvements, each validated by the branch comparisons:

• GCV smoothing (make_smoothing_spline, scipy ≥ 1.10) as default — eliminates the leading-arm under-smoothing wiggle. smoothing=<float> still available for explicit s.
• 6D closest-point matching (position + velocity) — more robust at orbital self-intersections. No cost for cold streams.
• Auto-timerange — essential for warm streams, harmless for cold.
• Boundary pile-up exclusion — removes particles matched to the far edge of the progenitor arc before fitting.

Warm-stream deep dive

After initially thinking the warm stream's track broke in 3D projections, discovered it was a plotting bug (pyplot.sca(ax) missing). The track is correct in all projections — verified numerically (74–436 sample particles within 3 kpc of each track evaluation point).

Stream-orbit refinement exploration (alt-stream-orbit)

Explored using a "stream orbit" integrated from the stream tip as the smoothing reference instead of the progenitor orbit:

• v1 (one-shot from raw particle-mean tip IC): diverges — noisy IC amplified over integration. Track flies outside sample cloud.
• v2 (two-step: fit preliminary track on progenitor → evaluate smooth track at tip for noise-free IC → integrate stream orbit from that): works — avoids divergence, matches main-branch quality.
• Conclusion: v2 reliably avoids v1's noise problem, but for the streams tested (LogarithmicHaloPotential and TriaxialNFW b=0.8, c=0.6, tdisrupt=5 Gyr), the main branch (progenitor orbit only) already works well. The refinement doesn't visibly improve the track. Left as experimental on alt-stream-orbit.

Triaxial NFW comparison

Ran main vs v1 vs v2 in TriaxialNFWPotential(b=0.8, c=0.6), tdisrupt=5 Gyr, for both cold (10^4 Msun) and warm (10^7 Msun) streams. Main handles both cleanly; v1 diverges for warm; v2 matches main. See earlier PR comments for plots.

Timing (track construction only, excluding particle sampling)

	n=1000	n=3000
Main	~1.0s	~2.3s
v2 (stream-orbit)	~2.0s	~4.6s

v2 is 2x main (fits the track twice). Progenitor integration is negligible (<60ms). Cost dominated by the N×M distance matrix and GCV spline fits. Potential choice and progenitor mass don't affect timing.

Current state

• streamspray-track: GCV + 6D + auto-timerange + boundary exclusion + percentile trim. 23 tests passing.
• alt-stream-orbit: experimental v2 refinement. Works but doesn't improve over main for tested potentials.
• 7 other alternative branches available for reference.

[jobovy]

Timing benchmarks (current implementation)

All times are track construction only — particle sampling/orbit integration excluded. The particles= argument lets users pre-sample once and reuse; sampling cost scales linearly with n and depends on the potential (e.g., ~17s for n=3000 in LogarithmicHalo, ~170s in TriaxialNFW).

scipy 1.17.1, make_smoothing_spline (GCV) for auto-smoothing.

LogarithmicHaloPotential (q=0.9), tdisrupt=4.5 Gyr

n	order=1 (mean)	order=2 (mean+cov)	reuse-s (order=2)
1000	0.13s	0.50s	0.09s
3000	0.22s	0.75s	0.17s
5000	0.33s	0.98s	0.24s

Cold (10⁴ M☉) and warm (10⁷ M☉) streams give essentially identical timing — progenitor mass doesn't affect track construction cost.

TriaxialNFWPotential (b=0.8, c=0.6), tdisrupt=5 Gyr

n	order=1 (mean)	order=2 (mean+cov)	reuse-s (order=2)
1000	0.19s	0.52s	0.14s
3000	0.28s	0.79s	0.22s

Slightly higher than LogarithmicHalo due to the longer progenitor arc (5 vs 4.5 Gyr), but same order of magnitude.

Cost breakdown

• GCV spline fitting dominates: order=2 fits 27 splines (6 mean + 21 unique covariance elements) vs. 6 for order=1, explaining the ~3x difference.
• KDTree matching: negligible (<50ms for n=3000). This replaced an N×M distance matrix that was ~1.6s for n=3000.
• Progenitor integration: <1ms (already integrated during spraydf setup).
• smoothing_s reuse: bypasses GCV entirely, using UnivariateSpline(s=...) with the stored per-spline smoothing values. Gives 3–4x speedup over GCV for order=2. Useful when re-fitting with different particles but similar stream properties.

Iteration (niter)

Each iteration re-matches particles to the current smooth track and refits. Cost is ~0.1s per iteration (n=3000, order=1):

niter	time (order=1)
0	0.21s
1	0.31s
2	0.41s
3	0.51s

Summary

Track construction is fast (sub-second for typical use). The dominant cost is always the upstream particle sampling, which scales with n and potential complexity. The particles= interface and smoothing_s reuse are the two main levers for avoiding redundant work.