Skip to content

PyJetscape: background/jet pair writer and the two-stage MUSIC production (prod_AuAu_0_10_jet) - #7

Merged
jhputschke merged 14 commits into
mainfrom
pair_h5_music
Sep 25, 2026
Merged

PyJetscape: background/jet pair writer and the two-stage MUSIC production (prod_AuAu_0_10_jet)#7
jhputschke merged 14 commits into
mainfrom
pair_h5_music

Conversation

@jhputschke

@jhputschke jhputschke commented Sep 25, 2026

Copy link
Copy Markdown
Owner

Summary

This adds a two-stage MUSIC pair production to PyJetscape. The X-SCAPE hydro runs MUSIC_1 for the background, then Matter + LBT + CausalLiquefier, then MUSIC_2 with the jet deposition. Both legs are written to one file in FastHydro's pair layout: arr = jet leg, arr_bg = background, plus source/droplets, shower/ and diag/. It mirrors prod_AuAu_0_10 with a jet.

Writer and capture

  • fno_h5_writer.py: pair extensions (0c55de0).
    • add_evolution() adds a second evolution that always has arr's shape.
    • ragged() writes streamed per-event tables, and write_diag() writes per-event scalars.
    • Every τ-bearing dataset gets a tau_axis attribute. repad_to now grows arr and arr_bg together.
    • Files are stamped freezeout_convention_id = "frames_written".
    • Single-leg H5BulkWriter output is unchanged apart from the new attributes.
  • showers.py and liquefier_io.py (68366ff): the shower and droplet capture moved from FastHydro, so PyJetscape can use it without importing FastHydro. FastHydro re-exports it.
  • pair_h5.py: PairH5Writer (9d3c4eb).
    • attach(jetscape) after Init() finds the legs, and sets dump_hydro_only on the jet leg only.
    • Both legs are read before anything is written, so a bad leg skips the event instead of leaving half a pair.
    • Built-in checks: diag/frames_identical, droplet counts and energies, bg_id for reuse, and diag/{bg,jet}_hit_boundary.

Bindings (38d8532, 3bce62b): MpiMusic.set_dump_hydro_only, set_skip_surface and get_hit_grid_boundary.

Production (483bf9a): example/prod_AuAu_0_10_jet/.

  • The job XML is derived from the prod XML, plus PythiaGun, CausalLiquefier, Matter + LBT and a second hydro block.
  • run_prod_jet.py options: --hard pythia|pgun, --reuse N, --no-deposit (the null test) and --no-showers, with XML guards.
  • run_jobs.sh is a wrapper around prod_AuAu_0_10/run_jobs.sh, which now reads PROD_SCRIPT, TAG_PREFIX and PROD_OUTDIR from the environment; its defaults are unchanged.

FastHydro (5d8a5c5): PairBrowser follows the file's freeze-out convention, so it opens these pair files.

Docs

  • The PyJetscape README has a new section, "Background/Jet Pair Writer", covering usage, layout, checks and timing, plus future speed steps.
  • PLAN_pair_h5_music.md archives the plan and its status.

Results

Setup: X-SCAPE music4gpu_test + JETSCAPE/X-SCAPE#141, music4gpu on a GB10, 0–10% Au+Au, full grid.

check result
null test (--no-deposit) arr == arr_bg in all 95 frames
deposition, seed 1 (PythiaGun 50–70 GeV, 25 droplets) the legs separate at the first deposit (τ = 1.1); the wake starts at the droplets
CPU vs GPU wake correlation ≥ 0.99999
--reuse 3 bg_id = [0,0,0] and identical arr_bg; arr differs
cost 60 s per pair event after the droplet pruning in #141 (184 s before), 58 s null test, 16 GB peak

Tests, run at the head of the follow-up branch pair_h5_music_surface_off, which contains this PR:

  • FastHydro: 315 passed, 8 skipped.
  • Visualization: 34 passed.
  • PyJetscape: 58 passed, 1 skipped, and 1 failure that is already on main. It's unrelated to this PR, which changes neither the test nor the SurfaceFinder code or bindings.
    • The failing test is test_surface_finder.py::test_lattice_spacing_is_honoured.
    • Its synthetic T(τ) crosses T_sw exactly at τ = 4.8, which is a lattice time, so Cornelius double-counts the flat surface: area 41.7 instead of 32.
    • Moving T_sw by 0.1% gives exactly 32. The fix belongs in the test, separately.

Merge notes

🤖 Generated with Claude Code

https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC

jhputschke and others added 14 commits September 24, 2026 16:56
@jhputschke @claude
PyJetscape: plan for the two-stage MUSIC pair production
5221ea7 
PLAN_pair_h5_music.md: run the X-SCAPE two-stage hydro (MUSIC_1 background ->
Matter+LBT+CausalLiquefier -> MUSIC_2 with the jet deposition) on the
prod_AuAu_0_10 physics and write both legs as a pair in FastHydro's layout
(arr = jet leg, arr_bg, source/droplets, shower/) from PyJetscape's own HDF5
writer. Covers the branch layout across X-SCAPE, MUSIC4GPU and js-contrib,
the music4gpu jet-source port, the writer extensions, the production driver
and how each piece is verified.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape: pair extensions for the FNO4d HDF5 writer
0c55de0 
FnoH5Writer can now write FastHydro's pair layout, so a MUSIC two-stage run
can store its background and jet legs in one file:

- add_evolution(name, fo_suffix): a second evolution (arr_bg) shaped,
  chunked and compressed like arr, with its own ntau_freezeout/tau_freezeout.
  ensure_capacity grows every evolution together, so arr_bg always has arr's
  shape (PairBrowser and wake_pyvista.py subtract them).
- write_frame/set_event_meta/append_event take dataset=; only the primary
  "arr" call marks an event written.
- ragged(): per-event variable-length tables (source/droplets, shower/...)
  written on every append instead of held until close. A table created late
  starts with empty events; set_event_meta refuses an event a table missed.
- write_diag(): per-event scalars under diag/, fast_data's dtypes (float64
  with NaN, uint64 seeds, strings), written per event.
- tau_axis attribute on every tau-bearing dataset; repad_to finds datasets by
  it and falls back to a legacy list that now includes arr_bg, so repadding a
  pair file no longer grows arr without arr_bg.
- freezeout_convention_id="frames_written" and has_source=false root attrs.

Single-leg H5BulkWriter output is unchanged apart from the new attributes.
tests/test_pair_h5.py covers the additions; test_h5_bulk.py stays green.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape: shower capture and liquefier readers usable by any writer
68366ff 
The droplet and shower capture never depended on the fast solver, but it
lived in FastHydro, which PyJetscape cannot import (FastHydro imports
jetscape). Move it so the MUSIC pair writer can use it:

- jetscape/showers.py: moved from fasthydro/showers.py unchanged apart from
  a note in the docstring (git mv keeps the history). fasthydro/showers.py
  re-exports it, loading the file directly in the js-contrib tree so that
  `import fasthydro` still keeps pyjetscape_core and torch off the import
  path (importing the jetscape package would load both).
- jetscape/liquefier_io.py: droplets(liq) -> (M, 8) float64 as stored by
  LiquefierBase::add_hydro_sources, and liquefier_params(liq) -> the five
  kernel parameters plus c_diff and gamma_relax. No fast_data or torch.
  FastHydro's params_from_liquefier/droplets_from_liquefier now use them.

FastHydro's full suite passes (313 passed, 8 skipped); new liquefier_io tests
in tests/test_pair_h5.py.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape: bind MpiMusic set_dump_hydro_only / set_skip_surface
38d8532 
Every MUSIC instance reads the first <Hydro><MUSIC> block, so a two-stage
run cannot give its legs different dump_hydro_only settings in XML. The
background leg must fill bulk_info (Matter/LBT read it); the jet leg should
keep its native store for the writer. Bind the existing C++ setters
(MusicWrapper.h) so the jet leg can be switched after JetScape.Init(), when
InitializeHydro() has already read the XML.

Not compiled yet: pyjetscape_core builds inside build_gpu, which cannot
rebuild until MUSIC4GPU has the jet source slot that X-SCAPE PR #138 calls.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape: PairH5Writer for two-stage MUSIC runs
9d3c4eb 
jetscape.pair_h5.PairH5Writer writes an X-SCAPE two-stage run (MUSIC_1
background -> Matter+LBT+CausalLiquefier -> MUSIC_2 with the deposition) as
one FNO4d-schema file in FastHydro's pair layout: arr = jet leg, arr_bg =
background, source/droplets (no source/S: MUSIC keeps no gridded source),
shower/, diag/, and the pair provenance and liquefier_* attributes.

- attach(jetscape) after Init(): finds the legs by module id, sets
  dump_hydro_only on the jet leg only (the background must keep filling
  bulk_info for Matter/LBT), finds the liquefier and the shower manager.
- The background is read from the framework copy, the jet leg from its
  native store; both go through the same resample onto one output grid.
  Both legs are read before anything is written, so an unreadable leg skips
  the event instead of leaving half a pair.
- diag/frames_identical counts the leading bit-identical frames. The writer
  warns when it is 0 (different initial conditions) or when a jet leg that
  received droplets is identical to the background (MUSIC ignored the
  liquefier, e.g. music4gpu without the jet source slot).
- diag records droplet counts/energies, including those depositing after
  the jet leg froze out (MUSIC never applies them), shower counts, MUSIC
  tau0, per-leg frame counts and bg_id, which repeats under setReuseHydro.
- Null test: with the jet leg's liquefier off, the writer clears the
  droplets itself.

Driven by hand between ExecPerEvent() and ClearPerEvent(), like
H5BulkWriter in run_prod.py; needs no compiled extension, so
tests/test_pair_h5.py runs it against stub legs.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape: prod_AuAu_0_10_jet, the two-stage MUSIC pair production
483bf9a 
example/prod_AuAu_0_10_jet/ mirrors prod_AuAu_0_10 with a jet: the same
0-10% Au+Au 3D MC-Glauber + MUSIC medium through the X-SCAPE two-stage hydro
(MUSIC_1 -> Matter+LBT+CausalLiquefier -> MUSIC_2), written by PairH5Writer.

- AuAu_MCGlauber_MUSIC_0_10_jet.xml: generated from the prod XML by exact
  substitution (MCGlauber, IS grid, PreEq and MUSIC physics unchanged), plus
  PythiaGun, CausalLiquefier (dtau = MUSIC Delta_Tau 0.02) and Matter+LBT
  from FastHydro's AuAu_FastHydro_tune_0_10_wake.xml, and a second Hydro
  block naming MUSIC_2. MUSIC_1 keeps dump_hydro_only 0 for the energy loss.
  Every tag has a default in jetscape_main.xml.
- run_prod_jet.py: run_prod.py's CLI (reusing its grid YAML, box and env
  helpers) plus --hard pythia|pgun, --pthat-min/max, --pgun-pt, --reuse N,
  --no-deposit (null test) and --no-showers. The job XML always holds one
  hard process (the automatic task list would run both). Guards: two Hydro
  blocks named MUSIC_1/MUSIC_2, MUSIC_1 dump_hydro_only 0, evolutionInMemory
  0, Eloss and Liquefier present, liquefier dtau == music_input Delta_Tau.
- run_jobs.sh: a wrapper around ../prod_AuAu_0_10/run_jobs.sh, which now
  takes PROD_SCRIPT, TAG_PREFIX and PROD_OUTDIR from the environment (its
  defaults, and so its behaviour, are unchanged).
- README.md: layout, options, how the legs are read, checks before a
  campaign. States that the chain has not run on MUSIC yet: it needs the
  jet source slot (MUSIC cee9460 / PR #138 on the CPU, the MUSIC4GPU port on
  the GPU), and time/memory are unmeasured.

Checked with --dry-run (both hard processes, reuse, null test, the guards)
and a two-job run_jobs.sh dry run.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
FastHydro: PairBrowser reads MUSIC pair files from PyJetscape
5d8a5c5 
- browse.py: both leg views now follow the file's freeze-out convention.
  fast_data files (no id) keep live() = ntau_freezeout - 1; files stamped
  freezeout_convention_id = "frames_written" (PyJetscape's H5BulkWriter and
  PairH5Writer) give live() = ntau_freezeout. Done in a view subclass,
  because fast_data/viz.py is vendored and never patched.
- tests/test_music_pair_browser.py: a pair written by jetscape.pair_h5 with
  stub MUSIC legs opens in PairBrowser: frame 0 identical, per-leg and pair
  live(), diff == arr - arr_bg, showers read back, summary provenance.
- README: the MUSIC two-stage section now points at PairH5Writer and
  prod_AuAu_0_10_jet and says how its three open problems were settled.

Full suite: 315 passed, 8 skipped.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape: record the pair_h5_music status in the plan
99e7936 
Phases 3-7 and 2c are committed; the MUSIC4GPU port, the builds and the
integration checks wait on switching the MUSIC4GPU checkout and pushing
music4gpu_test, which need the user.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape: verified on MUSIC; grid-boundary flag, bg_id fix, README …
3bce62b 
…corrections

Ran the two-stage pair production on build_gpu with the MUSIC4GPU jet-source
port: null test bit-identical over all frames, deposition separates the legs
at the first deposit time, CPU (MUSIC_FORCE_CPU) and GPU wakes agree, reuse
works. Fixes and additions from those runs:

- bind MpiMusic.get_hit_grid_boundary() (X-SCAPE ca8dd84a). PairH5Writer
  stores diag/bg_hit_boundary and diag/jet_hit_boundary, H5BulkWriter
  diag/hit_grid_boundary, and both warn: MUSIC stops an event whose
  freeze-out surface reaches the grid edge, so its evolution is truncated.
- PairH5Writer bg_id: hash the whole background leg, not its first frame.
  On MUSIC's native grid the first frame is at tau0, before any string
  deposits, and all zero, so every event got bg_id 0.
- prod_AuAu_0_10_jet README: measured cost (58 s/event null test, 184 s with
  deposition, 16 GB peak); the same seed does NOT give the prod_AuAu_0_10
  event (the extra modules draw from the random stream first), so that
  check is dropped.
- tests for both writer fixes; PLAN status updated.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape plan: keep XSCAPE and the Kokkos branches separate
25c6899 
No merges between MUSIC4GPU XSCAPE and KoKKos-Port / XSCAPE-KoKKos, nor
between X-SCAPE KoKKos-Music-Port and music4gpu_test (user decision). The
get_music4gpu.sh pin needs the port commit pushed to GitHub first. Notes the
one trap: an X-SCAPE branch with PR #138 cannot build against an unported
Kokkos branch.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape plan: branches pushed, music4gpu pin in place
3e8362b 
X-SCAPE music4gpu_test and pair_h5_music and js-contrib pair_h5_music are
on GitHub; get_music4gpu.sh pins MUSIC4GPU b9cc8be, which still has to be
pushed.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape: jet pairs at 60 s/event after droplet pruning
311a6c4 
X-SCAPE 896e3d1c and MUSIC4GPU 3037be7 evaluate only the liquefier droplets
that can deposit in the current step: the seed-1 0-10% event drops from
184 s to 60 s with bit-identical output. README speed numbers and PLAN
status updated.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape README: document the background/jet pair writer and its ti…
ff0fe06 
…ming

New section "Background/Jet Pair Writer (pair_h5.py)": what a two-stage
MUSIC pair is, the MUSIC / X-SCAPE / XML requirements, usage (the
prod_AuAu_0_10_jet driver and the Python attach/Exec loop), the file
layout, how the legs are read and the built-in checks, and the measured
timing: 60 s per 0-10% Au+Au pair event after the per-step droplet pruning
(184 s before, 58 s without deposition, ~23-25 s hydro-only), with a
breakdown by phase and inside each MUSIC evolution.

Also: new files in the Contents table, h5py prerequisite covers pair_h5.py,
and H5BulkWriter's diag/hit_grid_boundary is noted.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
@jhputschke @claude
PyJetscape README: future steps for the pair production speed
a0cfe46 
- Start MUSIC_2 from a full-state MUSIC_1 snapshot taken just before
  tau_delay: the legs are bit-identical until the first deposit (1.04 fm/c
  measured), so this skips the ~4 s string deposition and ~30 steps
  (~4-5 s/event). Explains why the snapshot time is the tau_delay bound
  (the droplets exist only after MUSIC_1; the bound costs ~2 steps), why the
  stored evolution history cannot be used (ideal fields in float every 5th
  step; no shear tensor, bulk pressure or previous step), what the hook
  needs (~0.3 GB, including the GPU working copy) and how to check it.
- Other levers from the timing breakdown: writer resampling (12.2 s),
  medium-store copy (6.2 s), CPU freeze-out surface (~5.6 s/run), idle CPU
  source pass (~1.3 s/run), upstream CPU MUSIC lacking the jet hook.
- PLAN "Still open" updated (the push item is done).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CrjXPHRbBKFbRYS7cdbKC
This was referenced Sep 25, 2026
Merged
Merged
@jhputschke
jhputschke merged commit c60d997 into main Sep 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@jhputschke