Stitched restart output retains only the original run's config

[AdemVessell]

I noticed a small output-provenance gap while checking restart reproducibility. On 7fb2d811d2c605c5d9144e89a8b9ea3394962b90, a stitched restart retains the original file's config attribute, with no saved configuration for the resumed segment.

In a short circular-geometry run, I changed generic heating from 1 MW to 2 MW at the restart and extended t_final from 0.3 to 0.5 seconds. The stitched data contains the 2 MW continuation through 0.5 seconds, but its saved config still says 1 MW, t_final=0.3, and restart=null. With stitch=False, the new config is saved correctly. The stitched continuation's data matches the unstitched continuation.

This appears to come from concat_datatrees: xarray keeps the first dataset's attributes after simulation_output_to_xr has populated the new config. The output documentation describes the saved config as usable for inspection or rerunning, so the limitation could surprise someone working only from the stitched file.

Reproducer, run from an environment with the pinned TORAX source installed:

import copy
import json
from pathlib import Path
from tempfile import TemporaryDirectory

import torax
import xarray as xr

base = {
    'profile_conditions': {'n_e_right_bc': 0.5e20},
    'numerics': {'t_final': 0.3, 'fixed_dt': 0.05},
    'plasma_composition': {},
    'geometry': {'geometry_type': 'circular', 'n_rho': 8},
    'sources': {'generic_heat': {'P_total': 1e6}},
    'pedestal': {},
    'transport': {'core_transport_models': {
        'prescribed': {'model_name': 'prescribed'}}},
    'solver': {'solver_type': 'linear'},
    'time_step_calculator': {'calculator_type': 'fixed'},
}
with TemporaryDirectory() as directory:
    initial, _ = torax.run_simulation(torax.ToraxConfig.from_dict(base),
                                       progress_bar=False)
    path = Path(directory) / 'initial.nc'
    initial.to_netcdf(path)
    restart_time = float(initial.time.values[3])
    cfg = copy.deepcopy(base)
    cfg['numerics'].update(t_initial=restart_time, t_final=0.5)
    cfg['sources']['generic_heat']['P_total'] = 2e6
    resumed = {}
    for stitch in (False, True):
        cfg['restart'] = dict(filename=path, time=restart_time,
                              do_restart=True, stitch=stitch)
        resumed[stitch], _ = torax.run_simulation(
            torax.ToraxConfig.from_dict(cfg), progress_bar=False)
    plain, stitched = resumed[False], resumed[True]
    times = plain.time.values[plain.time.values > restart_time]
    xr.testing.assert_equal(plain.sel(time=times), stitched.sel(time=times))
    assert stitched.attrs['config'] == initial.attrs['config']
    assert json.loads(plain.attrs['config'])['numerics']['t_final'] == 0.5
    saved = json.loads(stitched.attrs['config'])
    print('stored t_final/restart:', saved['numerics']['t_final'], saved['restart'])
    print('continuation power:', stitched['scalars']['P_aux_generic_total']
          .sel(time=times).values)

Observed output: stored t_final/restart: 0.3 None; continuation power is 2,000,000 W at every new time. Tested on macOS arm64, Python 3.12.13, JAX 0.11.2 and xarray 2026.7.0. The three existing concat/config output tests also pass. This is a metadata observation, not evidence of incorrect numerical results.

Would you prefer documenting this limitation first, or preserving the resumed config alongside the original in an agreed format? I can contribute a focused test and small change once that direction is clear.

I used Codex to help inspect the code, build the reproducer, and run these checks.

[jcitrin]

Thank you for noticing this and for the very clear explanation and repro. I can confirm the repro.

I think there are 2 things to do.

1. Ensure that the "restart" config is properly saved as an attr , which is the minimal viable option to ensure reproducibility
2. Additionally save the old config from which the restart was done, as an additional attr. This makes it fully self-contained.

Are you able to sign the CLA? If so, then please go ahead and contribute the PR. If not, that's fine, and we can take it from here.

[AdemVessell]

Thanks @jcitrin
I appreciated the guidance.
So, I’ve signed the CLA and opened #2521 with the configuration fix and regression tests.